From b649eb478e5ed75709a725cd08c070b24b737065 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Fri, 31 Jul 2026 15:52:19 +0800 Subject: [PATCH 01/24] feat: adapter the x_flash_attention to a5 --- .../beam_search_group/op_host/CMakeLists.txt | 25 + xllm_ops/build_aclnn.sh | 2 +- .../op_host/CMakeLists.txt | 18 +- .../op_kernel/x_attention_catlass_kernel.h | 1 - .../op_host/CMakeLists.txt | 43 +- .../a5_x_flash_attention_infer_tiling.h | 337 +++++++++ .../x_flash_attention_infer_tiling.cpp | 64 ++ .../op_host/x_flash_attention_infer_tiling.h | 16 + .../op_host/xfa_arch_config.h.in | 21 + .../arch35/a5_x_flash_attention_infer.h | 93 +++ .../a5_x_flash_attention_infer_kernel.h | 650 ++++++++++++++++++ .../a5_x_flash_attention_infer_kernel_utils.h | 233 +++++++ ..._x_flash_attention_infer_tiling_data_def.h | 59 ++ .../op_kernel/x_flash_attention_infer.cpp | 64 ++ 14 files changed, 1622 insertions(+), 4 deletions(-) create mode 100644 xllm_ops/x_flash_attention_infer/op_host/arch35/a5_x_flash_attention_infer_tiling.h create mode 100644 xllm_ops/x_flash_attention_infer/op_host/xfa_arch_config.h.in create mode 100644 xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer.h create mode 100644 xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer_kernel.h create mode 100644 xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer_kernel_utils.h create mode 100644 xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer_tiling_data_def.h diff --git a/xllm_ops/beam_search_group/op_host/CMakeLists.txt b/xllm_ops/beam_search_group/op_host/CMakeLists.txt index 9b6dbab..1ada90b 100644 --- a/xllm_ops/beam_search_group/op_host/CMakeLists.txt +++ b/xllm_ops/beam_search_group/op_host/CMakeLists.txt @@ -15,11 +15,36 @@ if (BUILD_OPEN_PROJECT) ) endif() +# Dynamically set CATLASS_ARCH based on the SOC being built. +# beam_search_group is a pure-vector op, but common/common.h unconditionally +# #includes "catlass/gemm/tile/tile_copy.hpp", whose dispatch headers +# (copy_gm_to_l1.hpp etc.) only bring in an implementation when CATLASS_ARCH is +# defined (2201=AtlasA2 / 3510=Ascend950). On the A5 (ascend950) build the macro +# was never injected, so the ascend950 cube templates were compiled out and the +# TileCopy aliases referred to a non-existent CopyGmToL1 -> "no template named". +# We follow the same SOC-aware injection as x_flash_attention_infer, but here we +# use the macro to explicitly distinguish A3 from A5 (per requirement): emit +# -DCATLASS_ARCH=3510 for ascend950/310p5 (A5, ascend950 cube path) +# -DCATLASS_ARCH=2201 otherwise (A3 ascend910_93 / A2 ascend910b, +# AtlasA2 cube path) +# Both branches MUST define CATLASS_ARCH, because tile_copy.hpp's dispatch header +# only brings in a CopyGmToL1 implementation when the macro is set. Leaving it +# empty would also break A3/A2 (CopyGmToL1 would not exist there either). +string(TOLOWER "${SOC_VERSION}" _BSG_SOC_LOWER) +string(TOLOWER "${ASCEND_COMPUTE_UNIT}" _BSG_UNIT_LOWER) +if(_BSG_SOC_LOWER MATCHES "ascend950" OR _BSG_SOC_LOWER MATCHES "ascend310p5" + OR _BSG_UNIT_LOWER MATCHES "ascend950" OR _BSG_UNIT_LOWER MATCHES "ascend310p5") + set(BSG_CATLASS_ARCH_DEF "-DCATLASS_ARCH=3510") # A5 (ascend950) +else() + set(BSG_CATLASS_ARCH_DEF "-DCATLASS_ARCH=2201") # A3 (ascend910_93) / A2 (ascend910b) +endif() + add_ops_compile_options( OP_NAME BeamSearchGroup OPTIONS --cce-auto-sync=on -Wno-deprecated-declarations -Werror + ${BSG_CATLASS_ARCH_DEF} -I${CANN_3RD_LIB_PATH}/catlass/include -I${CMAKE_CURRENT_LIST_DIR}/ -I${CMAKE_CURRENT_LIST_DIR}/../../../ diff --git a/xllm_ops/build_aclnn.sh b/xllm_ops/build_aclnn.sh index c08f339..9851ec6 100644 --- a/xllm_ops/build_aclnn.sh +++ b/xllm_ops/build_aclnn.sh @@ -329,7 +329,7 @@ elif [[ "$SOC_VERSION" =~ ^ascend950 ]]; then # ### JD's in-house operators #### "beam_search_group" - "x_attention" + # "x_attention" # A5 暂不参与编译, 集中解决 x_flash_attention_infer "cache_unshared_kv" "causal_conv1d" "causal_conv1d_qkv" diff --git a/xllm_ops/mc2/dispatch_ffn_combine/op_host/CMakeLists.txt b/xllm_ops/mc2/dispatch_ffn_combine/op_host/CMakeLists.txt index 72b8b2b..5330ce2 100644 --- a/xllm_ops/mc2/dispatch_ffn_combine/op_host/CMakeLists.txt +++ b/xllm_ops/mc2/dispatch_ffn_combine/op_host/CMakeLists.txt @@ -21,6 +21,22 @@ if (BUILD_OPEN_PROJECT) ) endif() +# Dynamically set CATLASS_ARCH based on the SOC being built (mirrors +# x_flash_attention_infer/op_host/CMakeLists.txt). In the CMake scope +# SOC_VERSION may be empty; the reliable variable is ASCEND_COMPUTE_UNIT +# (e.g. "ascend950"). We accept both spellings and any *950 / *310p5 variant, +# and inject -DCATLASS_ARCH=3510 for the A5(arch35) build; all other SOCs +# (AtlasA2/A3) fall back to -DCATLASS_ARCH=2201 so the catlass forwarding +# headers can still dispatch to the correct specialization. +string(TOLOWER "${SOC_VERSION}" _DFFN_SOC_LOWER) +string(TOLOWER "${ASCEND_COMPUTE_UNIT}" _DFFN_UNIT_LOWER) +if(_DFFN_SOC_LOWER MATCHES "ascend950" OR _DFFN_SOC_LOWER MATCHES "ascend310p5" + OR _DFFN_UNIT_LOWER MATCHES "ascend950" OR _DFFN_UNIT_LOWER MATCHES "ascend310p5") + set(CATLASS_ARCH_DEF "-DCATLASS_ARCH=3510") +else() + set(CATLASS_ARCH_DEF "-DCATLASS_ARCH=2201") +endif() + add_ops_compile_options( OP_NAME DispatchFFNCombine OPTIONS @@ -28,7 +44,7 @@ add_ops_compile_options( -Wno-deprecated-declarations -Werror -DHCCL_COMM - -DCATLASS_ARCH=2201 + ${CATLASS_ARCH_DEF} ${_DISPATCH_FFN_INC_OPTS} -I${CANN_3RD_LIB_PATH}/catlass/include ) diff --git a/xllm_ops/x_attention/op_kernel/x_attention_catlass_kernel.h b/xllm_ops/x_attention/op_kernel/x_attention_catlass_kernel.h index 1003d71..78bb7ec 100644 --- a/xllm_ops/x_attention/op_kernel/x_attention_catlass_kernel.h +++ b/xllm_ops/x_attention/op_kernel/x_attention_catlass_kernel.h @@ -20,7 +20,6 @@ limitations under the License. #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 "catlass/epilogue/dispatch_policy.hpp" #include "catlass/gemm/block/block_mmad.hpp" diff --git a/xllm_ops/x_flash_attention_infer/op_host/CMakeLists.txt b/xllm_ops/x_flash_attention_infer/op_host/CMakeLists.txt index 0c52a76..11b5881 100644 --- a/xllm_ops/x_flash_attention_infer/op_host/CMakeLists.txt +++ b/xllm_ops/x_flash_attention_infer/op_host/CMakeLists.txt @@ -4,7 +4,7 @@ # CANN Open Software License Agreement Version 2.0 (the "License"). # Please refer to the License for details. You may not use this file except in compliance with the License. # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# INCLUDING BUT NOT LIMITED TO NON-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() @@ -15,14 +15,55 @@ if (BUILD_OPEN_PROJECT) ) endif() +# Dynamically set CATLASS_ARCH based on the SOC being built. +# NOTE: In the CMake scope SOC_VERSION may be empty; the reliable variable is +# ASCEND_COMPUTE_UNIT (see CMakeCache, e.g. "ascend950"). We accept both and also +# any *950 / *310p5 spelling. The previous `SOC_VERSION STREQUAL "Ascend950"` +# check never matched, so -DCATLASS_ARCH=3510 was never injected and the A5 +# (arch35) branch was silently compiled out on the HOST side, causing the host +# tiling to fall back to XFAInferTilingData(SaveToBuffer) while the kernel read +# the 896B FATilingData -> field mis-alignment -> garbage tiling -> out all-zero. +string(TOLOWER "${SOC_VERSION}" _XFA_SOC_LOWER) +string(TOLOWER "${ASCEND_COMPUTE_UNIT}" _XFA_UNIT_LOWER) +if(_XFA_SOC_LOWER MATCHES "ascend950" OR _XFA_SOC_LOWER MATCHES "ascend310p5" + OR _XFA_UNIT_LOWER MATCHES "ascend950" OR _XFA_UNIT_LOWER MATCHES "ascend310p5") + set(CATLASS_ARCH_DEF "-DCATLASS_ARCH=3510") + set(_XFA_IS_A5 TRUE) +else() + set(CATLASS_ARCH_DEF "") + set(_XFA_IS_A5 FALSE) +endif() + add_ops_compile_options( OP_NAME XFlashAttentionInfer OPTIONS --cce-auto-sync=on -Wno-deprecated-declarations -Werror + ${CATLASS_ARCH_DEF} -I${CANN_3RD_LIB_PATH}/catlass/include ) +# CRITICAL: add_ops_compile_options only affects the op_impl(kernel) build, NOT +# the host tiling object (ophost_xllm_tiling_obj). That object library is defined +# by the CANN framework in another directory (beam_search/op_host), so neither +# set_source_files_properties nor add_compile_definitions from this subdir can +# reach it. Instead we generate a per-build config header next to the tiling +# sources; the header content depends on the SOC selected at CONFIGURE time, so +# the A3 build gets an empty header (falls back to XFAInferTilingData/SaveToBuffer) +# while the A5 build gets `#define CATLASS_ARCH 3510` (compiles the arch35 branch). +# The tiling sources include this header via a same-dir relative path, so it works +# regardless of which aggregated target actually compiles the .cpp. +if(_XFA_IS_A5) + set(XFA_ARCH_CONFIG_BODY "#define CATLASS_ARCH 3510") +else() + set(XFA_ARCH_CONFIG_BODY "") +endif() +configure_file( + ${CMAKE_CURRENT_SOURCE_DIR}/xfa_arch_config.h.in + ${CMAKE_CURRENT_SOURCE_DIR}/xfa_arch_config.h + @ONLY +) + if (NOT BUILD_OPS_RTY_KERNEL) add_modules_sources(OPTYPE x_flash_attention_infer ACLNNTYPE aclnn) endif() diff --git a/xllm_ops/x_flash_attention_infer/op_host/arch35/a5_x_flash_attention_infer_tiling.h b/xllm_ops/x_flash_attention_infer/op_host/arch35/a5_x_flash_attention_infer_tiling.h new file mode 100644 index 0000000..f26594c --- /dev/null +++ b/xllm_ops/x_flash_attention_infer/op_host/arch35/a5_x_flash_attention_infer_tiling.h @@ -0,0 +1,337 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-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 XLLM_OPS_XFAI_ARCH35_HOST_A5_X_FLASH_ATTENTION_INFER_TILING_H +#define XLLM_OPS_XFAI_ARCH35_HOST_A5_X_FLASH_ATTENTION_INFER_TILING_H + +#include +#include +#include +#include + +#include "../../op_kernel/arch35/a5_x_flash_attention_infer_tiling_data_def.h" + +namespace FAInferTiling { +constexpr int64_t SPARSE_MODE_INT_MAX = 2147483647; +constexpr int32_t SPARSE_MODE_NO_MASK = 0; +constexpr int32_t SPARSE_MODE_LEFT_UP = 1; +constexpr int32_t SPARSE_MODE_RIGHT_DOWN = 2; + +constexpr int32_t BLOCK_BASE_SIZE = 128; +constexpr uint32_t CV_RATIO = 2; +const int32_t WORKSPACE_BLOCK_SIZE_DB = 131072; + +struct FAInfo { + int64_t batchSize = 0; + int64_t numOfHeads = 0; + int64_t numOfKVHeads = 0; + int64_t seqSize = 0; + int64_t seqInnerSize = 0; + int64_t headSize = 0; + + uint32_t numBlocks = 0; + uint32_t blockSize = 0; + uint32_t maxBlockNumPerBatch = 0; + + uint32_t maskType = SPARSE_MODE_NO_MASK; + float scaleValue = 1.0; + int64_t* actualSeqLengths{nullptr}; + int64_t* actualSeqLengthsKV{nullptr}; +}; + +template +auto CeilDivision(T num1, T num2) -> T +{ + if (num2 == 0) { + return 0; + } + return (num1 + num2 - 1) / num2; +} + +template +auto CalcTailSize(T num1, T num2) -> T +{ + if (num2 == 0) { + return 0; + } + T mod = num1 % num2; + return mod != 0 ? mod : num2; +} + +inline void GetPreNextTokensLeftUp( + FATilingData& tilingData, int64_t actualSeqLength, int64_t actualSeqLengthKV, int64_t& preTokensLeftUp, + int64_t& nextTokensLeftUp) +{ + auto& baseParams = tilingData.inputParamsRegbase; + int64_t preTokens = SPARSE_MODE_INT_MAX; + int64_t nextTokens = SPARSE_MODE_INT_MAX; + if (baseParams.attenMaskCompressMode == SPARSE_MODE_LEFT_UP) { + preTokens = SPARSE_MODE_INT_MAX; + nextTokens = 0; + } + if (baseParams.attenMaskCompressMode == SPARSE_MODE_RIGHT_DOWN) { + preTokensLeftUp = SPARSE_MODE_INT_MAX; + nextTokensLeftUp = actualSeqLengthKV - actualSeqLength; + } else { + preTokensLeftUp = preTokens; + nextTokensLeftUp = nextTokens; + } +} + +inline void FixParamWithRowInvalid( + int64_t& actualSeqLength, int64_t actualSeqLengthKV, int64_t& preTokensLeftUp, int64_t& nextTokensLeftUp) +{ + int64_t nextTokensError = (nextTokensLeftUp < 0) ? -nextTokensLeftUp : 0; + int64_t preTokensError = (actualSeqLength > actualSeqLengthKV + preTokensLeftUp) ? + (actualSeqLength - actualSeqLengthKV - preTokensLeftUp) : + 0; + nextTokensLeftUp += nextTokensError; + preTokensLeftUp -= nextTokensError; + actualSeqLength -= nextTokensError; + actualSeqLength -= preTokensError; +} + +inline int64_t GetCutBlockNums( + int64_t blockSeqLengthKV, int64_t blockSeqLength, int64_t sInner, int64_t sOuter, int64_t token) +{ + if (sInner == 0 || sOuter == 0) { + return 0; + } + int64_t blockNums = 0; + int64_t blockToken = token > 0 ? ((token + sInner - 1) / sInner * sInner) : (token / sInner * sInner); + int64_t outDivIn = sOuter > sInner ? sOuter / sInner : 1; + int64_t InDivOut = sInner > sOuter ? sInner / sOuter : 1; + int64_t tolerance = 0; + int64_t smallSize = 0; + if (outDivIn >= 1) { + tolerance = outDivIn; + smallSize = sInner; + } else { + tolerance = InDivOut; + smallSize = sOuter; + } + int64_t innerCutBlockNums = (blockSeqLengthKV - blockToken) / smallSize - tolerance; + int64_t innerCutBlockLeftNums = -blockToken / smallSize - tolerance; + int64_t innerCutBlockDownNums = (blockSeqLengthKV - blockSeqLength - blockToken) / smallSize - tolerance; + int64_t tmpInnerCutBlockNums = + (innerCutBlockNums > 0) ? + (innerCutBlockNums % tolerance + innerCutBlockNums) * (innerCutBlockNums / tolerance + 1) / 2 : + 0; + blockNums += tmpInnerCutBlockNums; + int64_t tmpInnerCutBlockLeftNums = + (innerCutBlockLeftNums > 0) ? + (innerCutBlockLeftNums % tolerance + innerCutBlockLeftNums) * (innerCutBlockLeftNums / tolerance + 1) / 2 : + 0; + blockNums -= tmpInnerCutBlockLeftNums; + int64_t tmpInnerCutBlockDownNums = + (innerCutBlockDownNums > 0) ? + (innerCutBlockDownNums % tolerance + innerCutBlockDownNums) * (innerCutBlockDownNums / tolerance + 1) / 2 : + 0; + blockNums -= tmpInnerCutBlockDownNums; + return blockNums; +} + +inline int64_t GetCalcBlockNumsOneHead( + int64_t actualSeqLength, int64_t actualSeqLengthKV, int64_t sOuterSize, int64_t sInnerSize, int64_t preTokensLeftUp, + int64_t nextTokensLeftUp, bool isAttenMaskUsed) +{ + if (!isAttenMaskUsed) { + int64_t outerBlockNums = (actualSeqLength + sOuterSize - 1) / sOuterSize; + int64_t innerBlockNums = (actualSeqLengthKV + sInnerSize - 1) / sInnerSize; + int64_t toCalcBlockNums = innerBlockNums * outerBlockNums; + return toCalcBlockNums; + } else { + int64_t innerBlockNums = + (actualSeqLengthKV + static_cast(sInnerSize) - 1) / static_cast(sInnerSize); + int64_t blockSeqLengthKV = innerBlockNums * static_cast(sInnerSize); + int64_t outerBlockNums = + (actualSeqLength + static_cast(sOuterSize) - 1) / static_cast(sOuterSize); + int64_t blockSeqLength = outerBlockNums * static_cast(sOuterSize); + int64_t toCalcBlockNums = innerBlockNums * outerBlockNums; + toCalcBlockNums -= GetCutBlockNums( + blockSeqLengthKV, blockSeqLength, static_cast(sInnerSize), static_cast(sOuterSize), + nextTokensLeftUp); + toCalcBlockNums -= GetCutBlockNums( + blockSeqLengthKV, blockSeqLength, static_cast(sInnerSize), static_cast(sOuterSize), + blockSeqLengthKV - blockSeqLength + preTokensLeftUp); + return toCalcBlockNums; + } +} + +inline int64_t GetSInnerBlockNums(int64_t sInnerIndexStart, int64_t sInnerIndexEnd, int64_t innerBlockNums) +{ + int64_t sInnerBlockNums = 0; + if (sInnerIndexEnd < 0) { + sInnerBlockNums = 0; + } else if (sInnerIndexEnd < innerBlockNums) { + sInnerBlockNums = (sInnerIndexStart < 0) ? (sInnerIndexEnd + 1) : (sInnerIndexEnd - sInnerIndexStart + 1); + } else { + int64_t tmpSInnerBlockNums = sInnerIndexStart < innerBlockNums ? innerBlockNums - sInnerIndexStart : 0; + sInnerBlockNums = (sInnerIndexStart < 0) ? innerBlockNums : tmpSInnerBlockNums; + } + return sInnerBlockNums; +} + +// 对Batch/headNum/qSeqLen三根轴切多核策略,采用贪心切分,使得每个AI Core上的计算量尽可能均衡. +inline void ComputeSplitNBSeq( + FATilingData& tilingData, uint32_t batchSize, const size_t tilingElementArrayLen, + std::vector& actualSeqLengths, std::vector& actualSeqLengthsKV, int64_t sOuterSize, + int64_t sInnerSize, double coreWightTarget, uint32_t& curCore) +{ + auto& baseParams = tilingData.inputParamsRegbase; + std::vector bnAxisStartIdx(tilingElementArrayLen, 0U); + std::vector qSeqAxisStartIdx(tilingElementArrayLen, 0L); + int64_t curWeight = 0; + uint32_t lastHeadIdx = 0; // actual seq为0时不分配核 + uint32_t lastBatchIdx = 0; + uint32_t lastQSeqOuterIdx = 0; + for (uint32_t batchIdx = 0; batchIdx < batchSize; batchIdx++) { + for (uint32_t headNum = 0; headNum < baseParams.qHeads; headNum++) { + int64_t preTokensLeftUp = 0; + int64_t nextTokensLeftUp = 0; + GetPreNextTokensLeftUp( + tilingData, actualSeqLengths[batchIdx], actualSeqLengthsKV[batchIdx], preTokensLeftUp, + nextTokensLeftUp); + FixParamWithRowInvalid( + actualSeqLengths[batchIdx], actualSeqLengthsKV[batchIdx], preTokensLeftUp, nextTokensLeftUp); + int64_t outerBlockNums = (actualSeqLengths[batchIdx] + sOuterSize - 1) / sOuterSize; + int64_t innerBlockNums = (actualSeqLengthsKV[batchIdx] + sInnerSize - 1) / sInnerSize; + for (uint32_t sOuterIndex = 0; sOuterIndex < outerBlockNums; sOuterIndex++) { + int64_t diff = static_cast(coreWightTarget * double(curCore + 1)) - curWeight; + int64_t sInnerIndexStart = + -(preTokensLeftUp > 0 ? (preTokensLeftUp + sInnerSize - 1) / sInnerSize : + preTokensLeftUp / sInnerSize); + int64_t sInnerIndexEnd = nextTokensLeftUp > 0 ? (nextTokensLeftUp + sInnerSize - 1) / sInnerSize : + nextTokensLeftUp / sInnerSize; + int64_t sInnerBlockNums = GetSInnerBlockNums(sInnerIndexStart, sInnerIndexEnd, innerBlockNums); + if (sInnerBlockNums - diff > diff && + !(lastHeadIdx == 0 && lastBatchIdx == 0 && lastQSeqOuterIdx == 0)) { + curCore += 1; + bnAxisStartIdx[curCore] = batchIdx * baseParams.qHeads + headNum; + qSeqAxisStartIdx[curCore] = sOuterIndex; + } + lastHeadIdx = headNum + 1; + lastBatchIdx = batchIdx + 1; + lastQSeqOuterIdx = sOuterIndex + 1; + curWeight += sInnerBlockNums; + preTokensLeftUp -= sOuterSize; + nextTokensLeftUp += sOuterSize; + } + } + } + bnAxisStartIdx[curCore + 1] = batchSize * baseParams.qHeads; + qSeqAxisStartIdx[curCore + 1] = static_cast(lastQSeqOuterIdx); + + std::copy( + std::begin(bnAxisStartIdx), std::end(bnAxisStartIdx), + std::begin(tilingData.multiCoreParamsRegbase.bnAxisStartIdx)); + std::copy( + std::begin(qSeqAxisStartIdx), std::end(qSeqAxisStartIdx), + std::begin(tilingData.multiCoreParamsRegbase.sparseStartIdx)); +} + +inline void FillInputParams(const FAInfo& faInfo, FATilingData& tilingData) +{ + auto& inputParams = tilingData.inputParamsRegbase; + inputParams.batch = faInfo.batchSize; + inputParams.qHeads = faInfo.numOfHeads; + inputParams.kvHeads = faInfo.numOfKVHeads; + inputParams.groupSize = faInfo.numOfHeads / faInfo.numOfKVHeads; + inputParams.qSeqlen = faInfo.seqSize; + inputParams.kvSeqlen = faInfo.seqInnerSize; + inputParams.embed = faInfo.headSize; + inputParams.scaleValue = faInfo.scaleValue; + + inputParams.attenMaskCompressMode = faInfo.maskType; + inputParams.headNumRatio = static_cast(faInfo.numOfHeads / faInfo.numOfKVHeads); + inputParams.blockSize = faInfo.blockSize; + inputParams.blockTableDim2 = faInfo.maxBlockNumPerBatch; + inputParams.paBlockNumSum = faInfo.numBlocks; + inputParams.attenMaskQSeqlen = static_cast(faInfo.seqSize); + inputParams.attenMaskKvSeqlen = static_cast(faInfo.seqInnerSize); +} + +inline void FillActualSeqLengths( + const FAInfo& faInfo, FATilingData& tilingData, std::vector& actualSeqLengths, + std::vector& actualSeqLengthsKV) +{ + auto& inputParams = tilingData.inputParamsRegbase; + int64_t batchSize = inputParams.batch; + bool isActualSeqLengthsNull = (faInfo.actualSeqLengths == nullptr) ? true : false; + bool isActualSeqLengthsKVNull = (faInfo.actualSeqLengthsKV == nullptr) ? true : false; + auto actualSeqLengthsSize = (faInfo.actualSeqLengths == nullptr) ? batchSize : 0; + auto actualSeqLengthsKVSize = (faInfo.actualSeqLengthsKV == nullptr) ? batchSize : 0; + inputParams.isActualSeqLengthsNull = isActualSeqLengthsNull; + inputParams.isActualSeqLengthsKVNull = isActualSeqLengthsKVNull; + inputParams.actualSeqLengthsSize = static_cast(actualSeqLengthsSize); + inputParams.actualSeqLengthsKVSize = static_cast(actualSeqLengthsKVSize); + for (int64_t batchIdx = 0; batchIdx < batchSize; batchIdx++) { + if (isActualSeqLengthsNull) { + actualSeqLengths[batchIdx] = inputParams.qSeqlen; + } else { + actualSeqLengths[batchIdx] = faInfo.actualSeqLengths[batchIdx]; + } + if (isActualSeqLengthsKVNull) { + actualSeqLengthsKV[batchIdx] = inputParams.kvSeqlen; + } else { + actualSeqLengthsKV[batchIdx] = faInfo.actualSeqLengthsKV[batchIdx]; + } + } +} + +inline int32_t GetFATilingParam(const FAInfo& faInfo, uint32_t blockDim, FATilingData& faTilingData) +{ + FillInputParams(faInfo, faTilingData); + auto& inputParams = faTilingData.inputParamsRegbase; + int64_t batchSize = inputParams.batch; + std::vector actualSeqLengths(batchSize); + std::vector actualSeqLengthsKV(batchSize); + FillActualSeqLengths(faInfo, faTilingData, actualSeqLengths, actualSeqLengthsKV); + + bool isAttenMaskUsed = faInfo.maskType != SPARSE_MODE_NO_MASK; + int64_t totalBlockNumsOneHead = 0; + constexpr static auto sInnerSize = BLOCK_BASE_SIZE; + constexpr static auto sOuterSize = BLOCK_BASE_SIZE; + for (int64_t batchIdx = 0; batchIdx < batchSize; batchIdx++) { + int64_t actualSeqLengthsTmp = actualSeqLengths[batchIdx]; + int64_t preTokensLeftUp = 0; + int64_t nextTokensLeftUp = 0; + GetPreNextTokensLeftUp( + faTilingData, actualSeqLengths[batchIdx], actualSeqLengthsKV[batchIdx], preTokensLeftUp, nextTokensLeftUp); + FixParamWithRowInvalid(actualSeqLengthsTmp, actualSeqLengthsKV[batchIdx], preTokensLeftUp, nextTokensLeftUp); + totalBlockNumsOneHead += GetCalcBlockNumsOneHead( + actualSeqLengthsTmp, actualSeqLengthsKV[batchIdx], sOuterSize, sInnerSize, preTokensLeftUp, + nextTokensLeftUp, isAttenMaskUsed); + } + + double coreWeightTarget = (double(totalBlockNumsOneHead * inputParams.qHeads) / double(blockDim)); + int64_t qSeqlenOuterSize = (inputParams.qSeqlen + sOuterSize - 1) / sOuterSize; + const size_t tilingElementArrayLen = MAX_CORE_NUM; + uint32_t curIndx = 0; + ComputeSplitNBSeq( + faTilingData, batchSize, tilingElementArrayLen, actualSeqLengths, actualSeqLengthsKV, sOuterSize, sInnerSize, + coreWeightTarget, curIndx); + + int64_t sInnerBlockNum = (inputParams.kvSeqlen + sInnerSize - 1) / sInnerSize; + int64_t totalSize = (totalBlockNumsOneHead / sInnerBlockNum) * inputParams.qHeads; + + faTilingData.multiCoreParamsRegbase.qSeqlenOuterSize = qSeqlenOuterSize; + faTilingData.multiCoreParamsRegbase.coreNum = static_cast(curIndx + 1); + faTilingData.multiCoreParamsRegbase.totalSize = totalSize; + faTilingData.multiCoreParamsRegbase.splitFactorSize = CeilDivision(totalSize, static_cast(curIndx + 1)); + faTilingData.multiCoreParamsRegbase.splitFactorTailSize = + CalcTailSize(totalSize, faTilingData.multiCoreParamsRegbase.splitFactorSize); + + return 0; +} + +} // namespace FAInferTiling +#endif // XLLM_OPS_XFAI_ARCH35_HOST_A5_X_FLASH_ATTENTION_INFER_TILING_H \ No newline at end of file diff --git a/xllm_ops/x_flash_attention_infer/op_host/x_flash_attention_infer_tiling.cpp b/xllm_ops/x_flash_attention_infer/op_host/x_flash_attention_infer_tiling.cpp index 16f751a..b01579e 100644 --- a/xllm_ops/x_flash_attention_infer/op_host/x_flash_attention_infer_tiling.cpp +++ b/xllm_ops/x_flash_attention_infer/op_host/x_flash_attention_infer_tiling.cpp @@ -8,8 +8,13 @@  * See LICENSE in the root of the software repository for the full text of the License.  */ +// Per-build arch selection is provided via x_flash_attention_infer_tiling.h, +// which includes the CMake-generated xfa_arch_config.h before any struct def. #include #include "x_flash_attention_infer_tiling.h" +#if defined(CATLASS_ARCH) && (CATLASS_ARCH == 3510) +#include "arch35/a5_x_flash_attention_infer_tiling.h" +#endif #define ASCENDC_EXTERN_C namespace optiling { @@ -184,6 +189,9 @@ ge::graphStatus XFAInferTiling::RunTiling() FillSplitCoreTilingDataForJD(); SetWorkspaces(); +#if defined(CATLASS_ARCH) && (CATLASS_ARCH == 3510) + return RunTilingA5(); +#else // Save tilingData tiling_data_.SaveToBuffer(tiling_context_->GetRawTilingData()->GetData(), tiling_context_->GetRawTilingData()->GetCapacity()); @@ -191,8 +199,64 @@ ge::graphStatus XFAInferTiling::RunTiling() tiling_context_->SetBlockDim(cubeCoreNum); tiling_context_->SetTilingKey(GetTilingKey()); return ge::GRAPH_SUCCESS; +#endif } +#if defined(CATLASS_ARCH) && (CATLASS_ARCH == 3510) +ge::graphStatus XFAInferTiling::RunTilingA5() +{ + // 复用已解析到 tiling_data_ 的基础字段构造 FAInfo + int64_t batch = static_cast(tiling_data_.get_batch()); + int64_t qHeadNum = static_cast(tiling_data_.get_numHeads()); + int64_t kvHeadNum = static_cast(tiling_data_.get_kvHeads()); + int64_t embed = static_cast(tiling_data_.get_embeddingSize()); + int64_t numTokens = static_cast(tiling_data_.get_numTokens()); + uint32_t blockSize = tiling_data_.get_blockSize(); + uint32_t numBlocks = tiling_data_.get_numBlocks(); + uint32_t maxBlockNumPerBatch = tiling_data_.get_maxNumBlocksPerBatch(); + + // qSeqlen: 当前不支持不等长, 取平均 + int64_t qSeqlen = (batch > 0) ? (numTokens / batch) : numTokens; + // kvSeqlen: host 阶段无实际值, 用 paged cache 最大容量作默认 + int64_t kvSeqlen = static_cast(maxBlockNumPerBatch) * static_cast(blockSize); + + FAInferTiling::FAInfo faInfo{}; + faInfo.batchSize = batch; + faInfo.numOfHeads= qHeadNum; + faInfo.numOfKVHeads = kvHeadNum; + faInfo.seqSize = qSeqlen; + faInfo.seqInnerSize = kvSeqlen; + faInfo.headSize = embed; + faInfo.numBlocks = numBlocks; + faInfo.blockSize = blockSize; + faInfo.maxBlockNumPerBatch = maxBlockNumPerBatch; + faInfo.maskType = maskType; + faInfo.scaleValue = tiling_data_.get_scaleValue(); + faInfo.actualSeqLengths = nullptr; + faInfo.actualSeqLengthsKV = nullptr; + + FATilingData faTilingData{}; + int32_t ret2 = FAInferTiling::GetFATilingParam(faInfo, static_cast(cubeCoreNum), faTilingData); + if (ret2 != 0) { + return ge::GRAPH_FAILED; + } + int32_t coreNum = faTilingData.multiCoreParamsRegbase.coreNum; + if (coreNum <= 0) { + return ge::GRAPH_FAILED; + } + + auto rawTiling = tiling_context_->GetRawTilingData(); + if (rawTiling->GetCapacity() < sizeof(FATilingData)) { + return ge::GRAPH_FAILED; + } + std::memcpy(rawTiling->GetData(), &faTilingData, sizeof(FATilingData)); + rawTiling->SetDataSize(sizeof(FATilingData)); + tiling_context_->SetBlockDim(static_cast(coreNum)); + tiling_context_->SetTilingKey(GetTilingKey()); + return ge::GRAPH_SUCCESS; +} +#endif + ASCENDC_EXTERN_C ge::graphStatus TilingFunc(gert::TilingContext *context) { diff --git a/xllm_ops/x_flash_attention_infer/op_host/x_flash_attention_infer_tiling.h b/xllm_ops/x_flash_attention_infer/op_host/x_flash_attention_infer_tiling.h index 8aacd28..6e6da71 100644 --- a/xllm_ops/x_flash_attention_infer/op_host/x_flash_attention_infer_tiling.h +++ b/xllm_ops/x_flash_attention_infer/op_host/x_flash_attention_infer_tiling.h @@ -11,6 +11,12 @@ #ifndef __X_FLASH_ATTENTION_INFER_TILINGDATA_H__ #define __X_FLASH_ATTENTION_INFER_TILINGDATA_H__ +// Per-build arch selection (generated by CMake). On A5 this defines +// CATLASS_ARCH=3510 so the arch35 host tiling struct/branch is compiled in; on A3 +// it is empty. Included here (in the shared tiling header) so that EVERY +// translation unit that pulls in this header (tiling.cpp, proto.cpp, ...) sees a +// consistent XFAInferTilingData layout, avoiding ODR violations. +#include "xfa_arch_config.h" #include "register/tilingdata_base.h" #include "tiling/platform/platform_ascendc.h" #include "tiling/tiling_api.h" @@ -99,6 +105,13 @@ BEGIN_TILING_DATA_DEF(XFAInferTilingData) TILING_DATA_FIELD_DEF(uint64_t, splitOTotalSize); TILING_DATA_FIELD_DEF(uint32_t, totalSplitNodeNum); TILING_DATA_FIELD_DEF(uint32_t, needCoreNum); +#if defined(CATLASS_ARCH) && (CATLASS_ARCH == 3510) + // A5(arch35) path memcpy's the full 904B FATilingData into the RawTilingData + // buffer. The framework sizes RawTilingData from this registered struct, so we + // reserve a padding blob large enough to hold FATilingData. A3 builds do NOT + // define CATLASS_ARCH, so this field is absent and the A3 layout is unchanged. + TILING_DATA_FIELD_DEF_ARR(uint8_t, 1024, xfaA5TilingReserved); +#endif END_TILING_DATA_DEF; REGISTER_TILING_DATA_CLASS(XFlashAttentionInfer, XFAInferTilingData) @@ -115,6 +128,9 @@ class XFAInferTiling { ge::graphStatus FillBasicTilingData(); void FillSplitCoreTilingDataForJD(); void SetWorkspaces(); +#if defined(CATLASS_ARCH) && (CATLASS_ARCH == 3510) + ge::graphStatus RunTilingA5(); +#endif private: XFAInferTilingData tiling_data_; gert::TilingContext* tiling_context_ = nullptr; diff --git a/xllm_ops/x_flash_attention_infer/op_host/xfa_arch_config.h.in b/xllm_ops/x_flash_attention_infer/op_host/xfa_arch_config.h.in new file mode 100644 index 0000000..7d0d7be --- /dev/null +++ b/xllm_ops/x_flash_attention_infer/op_host/xfa_arch_config.h.in @@ -0,0 +1,21 @@ +// ----------------------------------------------------------------------------------------------------------- +// Copyright (c) 2025 Huawei Technologies Co., Ltd. +// This file is generated by CMake (configure_file) from xfa_arch_config.h.in. +// DO NOT EDIT the generated header directly. +// +// Purpose: The host tiling.cpp is compiled into a framework-aggregated OBJECT +// library (ophost_xllm_tiling_obj) defined in another directory, so per-source +// COMPILE_DEFINITIONS / add_compile_definitions from this op subdir cannot reach +// it. Instead we bake the A5 (arch35) selection into a generated header that the +// tiling sources include. The value is decided at CMake configure time based on +// ASCEND_COMPUTE_UNIT / SOC_VERSION, so A3 and A5 builds get different content. +// ----------------------------------------------------------------------------------------------------------- +#ifndef XFA_ARCH_CONFIG_H +#define XFA_ARCH_CONFIG_H + +// @XFA_ARCH_CONFIG_BODY@ is replaced by CMake: +// - on A5 (ascend950 / ascend310p5): "#define CATLASS_ARCH 3510" +// - otherwise : (empty) +@XFA_ARCH_CONFIG_BODY@ + +#endif // XFA_ARCH_CONFIG_H \ No newline at end of file diff --git a/xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer.h b/xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer.h new file mode 100644 index 0000000..a9f320d --- /dev/null +++ b/xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer.h @@ -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. + */ + +#ifndef XLLM_OPS_XFAI_ARCH35_A5_X_FLASH_ATTENTION_INFER_H +#define XLLM_OPS_XFAI_ARCH35_A5_X_FLASH_ATTENTION_INFER_H + +#include "a5_x_flash_attention_infer_kernel.h" + +namespace XllmOps { +namespace XfaArch35 { + +// A5(Ascend950/DAV_3510) dispatch helper. +// This is a device-callable (non-global) helper that inlines the example49 FAInferTla +// assembly logic, so the extern "C" global entry can dispatch into it directly. +// Parameter order follows the xllm_ops op_kernel entry signature. +template +CATLASS_DEVICE void FAInferA5Dispatch( + GM_ADDR query, GM_ADDR key_cache, GM_ADDR value_cache, GM_ADDR mask, GM_ADDR block_table, + GM_ADDR actual_q_lens, GM_ADDR actual_kv_lens, GM_ADDR attn_out, GM_ADDR tiling) +{ + using namespace Catlass; + using ArchTag = Arch::Ascend950; + using ElementQ = Dtype; + using LayoutTagQ = layout::RowMajor; + using ElementK = Dtype; + using LayoutTagK = layout::ColumnMajor; + using ElementV = Dtype; + using LayoutTagV = layout::RowMajor; + using ElementS = float; + using LayoutTagS = layout::RowMajor; + using ElementP = Dtype; + using LayoutTagP = layout::zN; + using ElementO = Dtype; + using LayoutTagO = layout::RowMajor; + using ElementMask = uint8_t; + using LayoutTagMask = layout::RowMajor; + using ElementOTmp = float; + using LayoutTagOTmp = layout::RowMajor; + // L1TileShape::K must be embedding + using L1TileShape = tla::Shape<_128, _128, _128>; + using L0TileShape = L1TileShape; + // GEMM Block: Flash Attention Infer Q * K^T + using DispatchPolicyQK = Gemm::MmadFAIQK; + using TileCopyQK = Gemm::Tile::PackedTileCopyTlaToUB< + ArchTag, ElementQ, LayoutTagQ, ElementK, LayoutTagK, ElementS, LayoutTagS, void, + Gemm::Tile::CopyL0CToUBMode::SPLIT_M>; + using TileMmadQK = Gemm::Tile::TileMmadTla; + using BlockMmadQK = Gemm::Block::BlockMmadTla< + DispatchPolicyQK, L1TileShape, L0TileShape, ElementQ, ElementK, ElementS, void, TileCopyQK, TileMmadQK>; + + // Epilogue Block: online softmax on current S base block + using DispatchPolicySoftmax = Epilogue::EpilogueAscend950FASoftmax; + using PType = Gemm::GemmType; + using SType = Gemm::GemmType; + using maskType = Gemm::GemmType; + using EpilogueOnlineSoftmax = + Epilogue::Block::BlockEpilogue; + + // GEMM Block: Flash Attention Infer P * V + using DispatchPolicyPV = Gemm::MmadFAIPV; + using TileCopyPV = Gemm::Tile::PackedTileCopyTlaToUB< + ArchTag, ElementP, LayoutTagP, ElementV, LayoutTagV, ElementOTmp, LayoutTagV, void, + Gemm::Tile::CopyL0CToUBMode::SPLIT_M>; + using TileMmadPV = Gemm::Tile::TileMmadTla; + using BlockMmadPV = Gemm::Block::BlockMmadTla< + DispatchPolicyPV, L1TileShape, L0TileShape, ElementP, ElementV, ElementOTmp, void, TileCopyPV, TileMmadPV>; + + // Epilogue Block: O base block rescale/update + using DispatchPolicyRescaleO = Epilogue::EpilogueAscend950FARescaleO; + using OType = Gemm::GemmType; + using OTmpType = Gemm::GemmType; + using EpilogueRescaleO = Epilogue::Block::BlockEpilogue; + + using FAInferKernelType = + FAInferKernel; + FAIKernelParams params{ + query, key_cache, value_cache, mask, block_table, actual_q_lens, actual_kv_lens, attn_out, tiling}; + // call kernel + FAInferKernelType flashAttnInfer; + flashAttnInfer(params); +} + +} // namespace XfaArch35 +} // namespace XllmOps + +#endif // XLLM_OPS_XFAI_ARCH35_A5_X_FLASH_ATTENTION_INFER_H \ No newline at end of file diff --git a/xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer_kernel.h b/xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer_kernel.h new file mode 100644 index 0000000..5db1da7 --- /dev/null +++ b/xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer_kernel.h @@ -0,0 +1,650 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-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 XLLM_OPS_XFAI_ARCH35_A5_X_FLASH_ATTENTION_INFER_KERNEL_H +#define XLLM_OPS_XFAI_ARCH35_A5_X_FLASH_ATTENTION_INFER_KERNEL_H + +// Device(kernel) side: the catlass forwarding headers (copy_gm_to_l1.hpp / copy_l1_to_l0a.hpp +// etc.) dispatch to the ascend950 specialization ONLY when CATLASS_ARCH == 3510. The kernel +// compile command does not inject -DCATLASS_ARCH (that is only injected on the host side by +// op_host/CMakeLists.txt). On A5 the device arch macro __NPU_ARCH__ == 3510, so derive +// CATLASS_ARCH from it here, BEFORE including any catlass header, so the Ascend950 tile-copy +// templates (CopyGmToL1 / CopyL1ToL0A / CopyL1ToL0B) are actually visible. +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) && !defined(CATLASS_ARCH) +#define CATLASS_ARCH 3510 +#endif + +#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 "catlass/epilogue/dispatch_policy.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" +#include "catlass/status.hpp" +#include "tla/layout.hpp" + +#include "kernel_operator.h" + +#include "a5_x_flash_attention_infer_kernel_utils.h" +#include "a5_x_flash_attention_infer_tiling_data_def.h" + +using namespace Catlass; +using namespace tla; +using namespace AscendC; + +template < + class BlockMmadQK, class BlockMmadPV, class EpilogueOnlineSoftmax, class EpilogueRescaleO, bool PAGED_CACHE_FLAG> +class FAInferKernel { +public: + using ArchTag = typename BlockMmadQK::ArchTag; + using L1TileShape = typename BlockMmadQK::L1TileShape; + using ElementQ = typename BlockMmadQK::ElementA; + using LayoutTagQ = typename BlockMmadQK::LayoutTagA; + using ElementK = typename BlockMmadQK::ElementB; + using LayoutTagK = typename BlockMmadQK::LayoutTagB; + using ElementS = typename BlockMmadQK::ElementC; + using LayoutTagS = typename BlockMmadQK::LayoutTagC; + + using ElementP = typename BlockMmadPV::ElementA; + using LayoutTagP = typename BlockMmadPV::LayoutTagA; + using ElementV = typename BlockMmadPV::ElementB; + using LayoutTagV = typename BlockMmadPV::LayoutTagB; + + using ElementMask = typename EpilogueOnlineSoftmax::ElementMask; + using LayoutTagMask = typename EpilogueOnlineSoftmax::LayoutTagMask; + + using ElementOTmp = typename EpilogueRescaleO::ElementOTmp; + using LayoutTagOTmp = typename EpilogueRescaleO::LayoutTagOTmp; + using ElementO = typename EpilogueRescaleO::ElementO; + using LayoutTagO = typename EpilogueRescaleO::LayoutTagO; + + static constexpr uint32_t qSeqlenTemplateType = tla::get<0>(L1TileShape{}); + static constexpr uint32_t kvSeqlenTemplateType = tla::get<1>(L1TileShape{}); + static constexpr uint32_t embedTemplateType = tla::get<2>(L1TileShape{}); + + static constexpr uint32_t MM2_LEFT_SIZE = qSeqlenTemplateType * kvSeqlenTemplateType * sizeof(ElementP); + + // Methods + CATLASS_DEVICE + FAInferKernel() + {} + + CATLASS_DEVICE void Init(FAIKernelParams const& params) + { + // 获取当前aic idx 和sub blockidx + if ASCEND_IS_AIC { + this->blockIdx = AscendC::GetBlockIdx(); + } else { + this->blockIdx = AscendC::GetBlockIdx() >> 1; + } + + this->subBlockIdx = AscendC::GetSubBlockIdx(); + constInfo.subBlockIdx = this->subBlockIdx; + + // 调用Tiling接口 + auto faTilingStruct = (__gm__ FATilingData*)params.tiling; + auto& inputParamsRegbase = faTilingStruct->inputParamsRegbase; + this->constInfo.scaleValue = static_cast(inputParamsRegbase.scaleValue); + this->constInfo.batch = inputParamsRegbase.batch; + this->constInfo.qHeads = inputParamsRegbase.qHeads; + this->constInfo.kvHeads = inputParamsRegbase.kvHeads; + this->constInfo.groupSize = inputParamsRegbase.groupSize; + this->constInfo.qSeqlen = inputParamsRegbase.qSeqlen; + this->constInfo.kvSeqlen = inputParamsRegbase.kvSeqlen; + this->constInfo.embed = inputParamsRegbase.embed; + this->constInfo.attenMaskQSeqlen = inputParamsRegbase.attenMaskQSeqlen; + this->constInfo.attenMaskKvSeqlen = inputParamsRegbase.attenMaskKvSeqlen; + + this->constInfo.headNumRatio = inputParamsRegbase.headNumRatio; + this->constInfo.actualSeqLengthsSize = inputParamsRegbase.actualSeqLengthsSize; + this->constInfo.actualSeqLengthsKVSize = inputParamsRegbase.actualSeqLengthsKVSize; + this->constInfo.isActualSeqLengthsNull = inputParamsRegbase.isActualSeqLengthsNull; + this->constInfo.isActualSeqLengthsKVNull= inputParamsRegbase.isActualSeqLengthsKVNull; + + // pageAttention + if constexpr (PAGED_CACHE_FLAG) { + this->constInfo.blockTableDim2 = inputParamsRegbase.blockTableDim2; + this->constInfo.blockSize = inputParamsRegbase.blockSize; + this->constInfo.paBlockNumSum = inputParamsRegbase.paBlockNumSum; + } + + auto& multiCoreParamsRegbase = faTilingStruct->multiCoreParamsRegbase; + this->constInfo.qSeqlenOuterSize = multiCoreParamsRegbase.qSeqlenOuterSize; + this->constInfo.coreNum = multiCoreParamsRegbase.coreNum; + /* 多核切分偏移计算 */ + this->constInfo.multiCoreInnerOffset = multiCoreParamsRegbase.sparseStartIdx[this->blockIdx]; + this->constInfo.multiCoreInnerLimit = multiCoreParamsRegbase.sparseStartIdx[this->blockIdx + 1]; + this->constInfo.bnAxisStartIdx = multiCoreParamsRegbase.bnAxisStartIdx[this->blockIdx]; + this->constInfo.bnAxisEndIdx = multiCoreParamsRegbase.bnAxisStartIdx[this->blockIdx + 1]; + + CrossCoreSetFlag(MM2_RES_INTRA_EVENT[0]); + CrossCoreSetFlag(MM2_RES_INTRA_EVENT[1]); + CrossCoreSetFlag(MM1_RES_INTRA_EVENT[0]); + CrossCoreSetFlag(MM1_RES_INTRA_EVENT[1]); + + this->constInfo.qSeqlenBase = qSeqlenTemplateType; + this->constInfo.kvSeqlenBase = kvSeqlenTemplateType; + + for (int i = 0; i < NUM2; i++) { + bmm1TensorList[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += MM1_RESULT_SIZE; + bmm2TensorList[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += MM2_RESULT_SIZE; + } + + if ASCEND_IS_AIV { + for (int i = 0; i < KERNEL_TASK_NUM; i++) { + sumUb[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += SHARE_UB_SIZE; + expUb[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += SHARE_UB_SIZE; + maxUb[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += SHARE_UB_SIZE; + } + } + + // 初始化全局L1 + for (int i = 0; i < KERNEL_TASK_NUM; i++) { + mm2AL1TensorList[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + i * MM2_LEFT_SIZE); + } + l1BufAddrStart += KERNEL_TASK_NUM * MM2_LEFT_SIZE; + } + + CATLASS_DEVICE void operator()(FAIKernelParams const& params) + { + // Init + Init(params); // 初始化ConstInfo，初始化L1/UB + uint32_t l0CBufAddrStart = 0; + BlockMmadQK blockMmadMmadQK(resource, l1BufAddrStart, l0CBufAddrStart); + BlockMmadPV blockMmadMmadPV(resource, l1BufAddrStart, l0CBufAddrStart); + EpilogueOnlineSoftmax epilogueOnlineSoftmax(resource, constInfo.scaleValue, ubBufAddrStart); + EpilogueRescaleO epilogueRescaleO(resource, ubBufAddrStart); + + // Get blockIdx + int32_t blockNum = this->constInfo.coreNum; + if (this->blockIdx >= blockNum) { + return; + } + + int64_t batch = this->constInfo.batch; + int64_t qHeads = this->constInfo.qHeads; + int64_t qSeqlen = this->constInfo.qSeqlen; + int64_t kvHeads = this->constInfo.kvHeads; + int64_t kvSeqlen = this->constInfo.kvSeqlen; + int64_t groupSize = this->constInfo.groupSize; + int64_t embed = this->constInfo.embed; + int64_t blockSize = this->constInfo.blockSize; + // Read runtime actual kv length from the actualSeqLengthsKV tensor so that the + // causal mask can shield the padding tail when actual_kv_len < paged-cache capacity. + // The tiling-provided constInfo.kvSeqlen is the paged-cache capacity (maxBlockNumPerBatch + // * blockSize), NOT the real kv length, so relying on it makes decode attend to padded kv + // (aligns with the A3 path which reads gActualKvseqlen.GetValue(BIdx)). + int64_t actualKvSeqlen = kvSeqlen; + if (params.actualKvSeqlen != nullptr) { + AscendC::GlobalTensor gActualKvseqlen; + gActualKvseqlen.SetGlobalBuffer((__gm__ int32_t*)params.actualKvSeqlen); + int64_t rtKv = static_cast(gActualKvseqlen.GetValue(0)); + if (rtKv > 0) { + actualKvSeqlen = rtKv; + } + } + int64_t kvSeqlenMask = actualKvSeqlen; + // causal mask uses a shared [MASK_DIM, MASK_DIM] triu template (row = query + // absolute logical position, col = kv absolute position), aligned with the A3 + // path which hardcodes LayoutMask(2048, 2048). diffS shifts single-token decode + // query to the causal bottom so it can attend to all valid kv. diffS MUST use the + // actual kv length (not the paged-cache capacity), otherwise the causal mask fails + // to shield the padded tail (e.g. actual_kv=64 within a 128-slot block). + constexpr int64_t MASK_DIM = 2048; + int64_t diffS = actualKvSeqlen - qSeqlen; + if constexpr (PAGED_CACHE_FLAG) { + kvSeqlen = RoundUp(kvSeqlen, blockSize); + } + // Init Tensor + AscendC::GlobalTensor gmQ; + gmQ.SetGlobalBuffer((__gm__ ElementQ*)params.q); + // Create TLA layouts for kernel usage + auto layoutQ = MakeLayout(batch * qSeqlen, kvHeads * groupSize * embed); + auto tensorQWithLayout = tla::MakeTensor(gmQ, layoutQ, Arch::PositionGM{}); + + AscendC::GlobalTensor gmK; + gmK.SetGlobalBuffer((__gm__ ElementK*)params.k); + auto layoutK = MakeLayout(kvHeads * embed, batch * kvSeqlen); + auto tensorKWithLayout = tla::MakeTensor(gmK, layoutK, Arch::PositionGM{}); + + AscendC::GlobalTensor gmV; + gmV.SetGlobalBuffer((__gm__ ElementV*)params.v); + auto layoutV = MakeLayout(batch * kvSeqlen, kvHeads * embed); + auto tensorVWithLayout = tla::MakeTensor(gmV, layoutV, Arch::PositionGM{}); + + AscendC::GlobalTensor gmMask; + gmMask.SetGlobalBuffer((__gm__ ElementMask*)params.mask); + auto layoutMask = MakeLayout(MASK_DIM, MASK_DIM); + auto tensorMaskWithLayout = tla::MakeTensor(gmMask, layoutMask, Arch::PositionGM{}); + + // BlockTable + AscendC::GlobalTensor tensorTable; + tensorTable.SetGlobalBuffer((__gm__ int32_t*)params.blockTables); + + AscendC::GlobalTensor attentionOutGm; + AscendC::GlobalTensor workspaceGm; + attentionOutGm.SetGlobalBuffer((__gm__ ElementO*)params.o); + auto layoutO = MakeLayout(batch * qSeqlen, kvHeads * groupSize * embed); + auto attentionOutGmWithLayout = tla::MakeTensor(attentionOutGm, layoutO, Arch::PositionGM{}); + + uint32_t maxBlockNumPerBatch = this->constInfo.blockTableDim2; + + // Main process loop + + // 确定核内切分起点 + int64_t qSeqAxisStartIdx; + uint32_t bnAxisStartIdx; + uint32_t bnAxisEndIdx; + int64_t kvSeqLoopLimit; + int64_t nextQSeqAxisIdx = this->constInfo.multiCoreInnerLimit; + bnAxisStartIdx = this->constInfo.bnAxisStartIdx; + qSeqAxisStartIdx = this->constInfo.multiCoreInnerOffset; + if (likely((this->constInfo.coreNum - 1) > this->blockIdx)) { + bnAxisEndIdx = this->constInfo.bnAxisEndIdx; + if (nextQSeqAxisIdx != 0) { + bnAxisEndIdx++; + } + } else { + bnAxisEndIdx = this->constInfo.batch * this->constInfo.kvHeads * this->constInfo.headNumRatio; + } + + // 初始化CV流水状态信息 + int64_t taskId = 0; + bool notLast = true; + bool isLastBmm1 = false; + int64_t multiCoreInnerIdx = 1; + for (uint32_t bnIdx = bnAxisStartIdx; bnIdx < bnAxisEndIdx; ++bnIdx) { + bool lastBN = (bnIdx == bnAxisEndIdx - 1); + runParam.batchOuterIdx = bnIdx / (this->constInfo.kvHeads * this->constInfo.headNumRatio); + runParam.kvHeadsOuterIdx = + (bnIdx / this->constInfo.headNumRatio) % this->constInfo.kvHeads; // 切核逻辑，先N2G再B + ComputeParamBatch(runParam, this->constInfo, this->attenMaskInfo); // 计算runParam中参数值 + ComputeQseqLoopInfo(runParam, this->constInfo, lastBN, nextQSeqAxisIdx); + int64_t tempQSeqAxisEnd = lastBN ? (runParam.qSeqLoopTimes + 3) : runParam.qSeqLoopTimes; + for (int64_t qSeqAxisIndex = qSeqAxisStartIdx; qSeqAxisIndex < tempQSeqAxisEnd; ++qSeqAxisIndex) { + bool notLastThreeLoop = true; + bool notLastTwoLoop = true; + if (lastBN) { + int32_t extraQSeqAxis = qSeqAxisIndex - runParam.qSeqLoopTimes; + switch (extraQSeqAxis) { + case -1: + isLastBmm1 = true; + break; + case 0: + notLastThreeLoop = false; + break; + case 1: + notLastThreeLoop = false; + notLastTwoLoop = false; + break; + case 2: + notLast = false; + notLastThreeLoop = false; + notLastTwoLoop = false; + break; + default: + break; + } + } + if (notLastThreeLoop) { + runParam.groupIdx = bnIdx % this->constInfo.headNumRatio; + runParam.qSeqOuterAxisIdx = qSeqAxisIndex % this->constInfo.qSeqlenOuterSize; + ComputeParamQSeq(runParam, this->constInfo, qSeqAxisIndex); + ComputeKvSeqLoopInfo(runParam, this->constInfo); + kvSeqLoopLimit = runParam.kvSeqLoopEndIdx - 1; + } else { + runParam.kvSeqLoopStartIdx = 0; + kvSeqLoopLimit = 0; + } + for (int64_t kvSeqLoopCount = runParam.kvSeqLoopStartIdx; kvSeqLoopCount <= kvSeqLoopLimit; + ++kvSeqLoopCount) { + if (notLastThreeLoop) { + RunInfo& runInfo1 = runInfo[taskId & 3]; + this->SetRunInfo(runInfo1, runParam, taskId, kvSeqLoopCount, kvSeqLoopLimit, multiCoreInnerIdx); + if ASCEND_IS_AIC { + CalcKvSeqCoord(runInfo1, this->constInfo); + CalcQSeqCoord(runInfo1, this->constInfo); + auto actualShape = + tla::MakeShape(runInfo1.qSeqRealSize, runInfo1.kvSeqRealSize, this->constInfo.embed); + auto layoutMM1O = + tla::MakeLayout(runInfo1.qSeqRealSize, kvSeqlenTemplateType); + auto tensorMM1OWithLayout = + tla::MakeTensor(bmm1TensorList[runInfo1.taskIdMod2], layoutMM1O, Arch::PositionUB{}); + + auto tensorInQ = GetTile( + tensorQWithLayout, + tla::MakeCoord( + runInfo1.batchOuterIdx * qSeqlen + coordInfo[runInfo1.taskIdMod3].qSeqCoord, + runInfo1.kvHeadsOuterIdx * groupSize * embed + runInfo1.groupIdx * embed), + tla::MakeShape(runInfo1.qSeqRealSize, this->constInfo.embed)); + auto kCoord = runInfo1.kvHeadsOuterIdx * embed; + auto nCoord = 0; + auto nShape = runInfo1.kvSeqRealSize; + if constexpr (PAGED_CACHE_FLAG) { + uint32_t maxBlockNumPerBatch = this->constInfo.blockTableDim2; + uint64_t blockTableBaseOffset = + runInfo1.batchOuterIdx * maxBlockNumPerBatch; // 块表的基偏移量 + uint32_t curKvSeqAxisIdx = runInfo1.kvSeqLoopCount * this->constInfo.kvSeqlenBase; + uint64_t blockIdOffset = + curKvSeqAxisIdx / this->constInfo.blockSize; // 获取block table上的索引 + runInfo1.blockTableOffset = blockTableBaseOffset + blockIdOffset; + nShape = batch * kvSeqlen; + } else { + nCoord = coordInfo[runInfo1.taskIdMod3].curBIdx * kvSeqlen + + coordInfo[runInfo1.taskIdMod3].kvSeqCoord; + } + auto tensorInK = GetTile( + tensorKWithLayout, tla::MakeCoord(kCoord, nCoord), + tla::MakeShape(this->constInfo.embed, nShape)); + + auto tensorInTable = tensorTable[runInfo1.blockTableOffset]; + + bool isFirstLoop = (runInfo1.kvSeqLoopCount == runInfo1.kvSeqLoopStartIdx) ? true : false; + bool isLastUpdate = (runInfo1.kvSeqLoopCount == runInfo1.kvSeqLoopLimit) ? true : false; + + blockMmadMmadQK( + tensorInQ, tensorInK, tensorMM1OWithLayout, tensorInTable, actualShape, + runInfo1.taskIdMod2, this->constInfo.blockSize, isFirstLoop, isLastUpdate); + + CrossCoreSetFlag( + SYNC_C1_V1_FLAG[runInfo1.taskIdMod2]); // fixpip将结果搬运到UB后，设置SYNC_C1_V1_FLAG + CrossCoreSetFlag( + 16 + + SYNC_C1_V1_FLAG[runInfo1.taskIdMod2]); // fixpip将结果搬运到UB后，设置SYNC_C1_V1_FLAG + } + } + + if (taskId > 0 && notLastTwoLoop) { + if ASCEND_IS_AIV { + auto& runInfo3 = runInfo[(taskId + 3) & 3]; + auto& taskIdMod2 = runInfo3.taskIdMod2; + auto& taskIdMod3 = runInfo3.taskIdMod3; + auto& multiCoreIdxMod3 = runInfo3.multiCoreIdxMod3; + bool isFirstLoop = (runInfo3.kvSeqLoopCount == runInfo3.kvSeqLoopStartIdx) ? true : false; + CrossCoreWaitFlag( + SYNC_C1_V1_FLAG[taskIdMod2]); // 等待bmm1完成/等待SYNC_C1_V1_FLAG置位 + auto bmm1Layout = tla::MakeLayout( + runInfo3.halfQSeqRealSize, runInfo3.kvSeqRealSize); + auto bmm1Tensor = + tla::MakeTensor(bmm1TensorList[taskIdMod2], bmm1Layout, Arch::PositionUB{}); + auto l1Vf1OutLayout = + tla::MakeLayout(qSeqlenTemplateType, kvSeqlenTemplateType); + auto l1Vf1OutTensor = + tla::MakeTensor(mm2AL1TensorList[taskIdMod3], l1Vf1OutLayout, Arch::PositionL1{}); + + auto l1Vf1OutTile = GetTile( + l1Vf1OutTensor, + tla::MakeCoord(constInfo.subBlockIdx * runInfo3.firstHalfQSeqRealSize, 0), + tla::MakeShape(runInfo3.halfQSeqRealSize, kvSeqlenTemplateType)); + + // mask is a shared [MASK_DIM, MASK_DIM] triu template: row = query + // absolute logical position within the sequence (NOT batch-offset, + // the template is shared across batches like the A3 path). diffS aligns + // the query row to the causal bottom so decode (qSeqlen=1) attends to all kv. + int64_t qSeqOffset = runInfo3.qSeqOuterAxisIdx * qSeqlenTemplateType + + runInfo3.firstHalfQSeqRealSize * constInfo.subBlockIdx; + int64_t kvSeqOffset = runInfo3.kvSeqLoopCount * kvSeqlenTemplateType; + + auto gmMaskTile = GetTile( + tensorMaskWithLayout, tla::MakeCoord(diffS + qSeqOffset, kvSeqOffset), + tla::MakeShape(runInfo3.halfQSeqRealSize, runInfo3.kvSeqRealSize)); + + epilogueOnlineSoftmax( + l1Vf1OutTile, sumUb[multiCoreIdxMod3], maxUb[multiCoreIdxMod3], expUb[taskIdMod3], + bmm1Tensor, gmMaskTile, !isFirstLoop, taskIdMod2, taskIdMod3, + MM1_RES_INTRA_EVENT[taskIdMod2], SYNC_V1_C2_FLAG[taskIdMod3]); + } + } + if (taskId > 1 && notLast) { + if ASCEND_IS_AIC { + RunInfo& runInfo2 = runInfo[(taskId + 2) & 3]; + auto& taskIdMod2 = runInfo2.taskIdMod2; + auto& taskIdMod3 = runInfo2.taskIdMod3; + CrossCoreWaitFlag(SYNC_V1_C2_FLAG[taskIdMod3]); + CrossCoreWaitFlag(16 + SYNC_V1_C2_FLAG[taskIdMod3]); + + auto layoutMM2O = + tla::MakeLayout(runInfo2.qSeqRealSize, embedTemplateType); + auto mm2OutTensor = + tla::MakeTensor(bmm2TensorList[taskIdMod2], layoutMM2O, Arch::PositionUB{}); + + auto layoutVec1O = + tla::MakeLayout(qSeqlenTemplateType, kvSeqlenTemplateType); + auto mm2AL1Tensor = + tla::MakeTensor(mm2AL1TensorList[taskIdMod3], layoutVec1O, Arch::PositionL1{}); + auto kCoord = 0; + auto nCoord = runInfo2.kvHeadsOuterIdx * embed; + auto kShape = runInfo2.kvSeqRealSize; + if constexpr (PAGED_CACHE_FLAG) { + kShape = batch * kvSeqlen; + } else { + kCoord = coordInfo[runInfo2.taskIdMod3].curBIdx * kvSeqlen + + coordInfo[runInfo2.taskIdMod3].kvSeqCoord; + } + auto tensorInV = GetTile( + tensorVWithLayout, tla::MakeCoord(kCoord, nCoord), + tla::MakeShape(kShape, this->constInfo.embed)); + auto actualShape = + tla::MakeShape(runInfo2.qSeqRealSize, embedTemplateType, runInfo2.kvSeqRealSize); + auto tensorInTableV = tensorTable[runInfo2.blockTableOffset]; + blockMmadMmadPV( + mm2AL1Tensor, tensorInV, mm2OutTensor, tensorInTableV, actualShape, taskIdMod2, + this->constInfo.blockSize); + CrossCoreSetFlag( + SYNC_C2_V2_FLAG[runInfo2.taskIdMod2]); // fixpip将结果搬运到UB后，设置SYNC_C2_V2_FLAG + CrossCoreSetFlag( + 16 + + SYNC_C2_V2_FLAG[runInfo2.taskIdMod2]); // fixpip将结果搬运到UB后，设置SYNC_C2_V2_FLAG + } + } + if (taskId > 2) { + if ASCEND_IS_AIV { + RunInfo& runInfo3 = runInfo[(taskId + 1) & 3]; + auto& taskIdMod2 = runInfo3.taskIdMod2; + auto& taskIdMod3 = runInfo3.taskIdMod3; + auto& multiCoreIdxMod3 = runInfo3.multiCoreIdxMod3; + + bool isFirstLoop = (runInfo3.kvSeqLoopCount == runInfo3.kvSeqLoopStartIdx) ? true : false; + bool isLastUpdate = (runInfo3.kvSeqLoopCount == runInfo3.kvSeqLoopLimit) ? true : false; + CrossCoreWaitFlag( + SYNC_C2_V2_FLAG[taskIdMod2]); // 等待bmm2完成/等待SYNC_C2_V2_FLAG置位 + auto bmm2Layout = + MakeLayout(runInfo3.halfQSeqRealSize, embedTemplateType); + auto bmm2Tensor = + tla::MakeTensor(bmm2TensorList[taskIdMod2], bmm2Layout, Arch::PositionUB{}); + int64_t bOffset = runInfo3.batchOuterIdx * qSeqlen; + int64_t qSeqOffset = runInfo3.qSeqOuterAxisIdx * qSeqlenTemplateType + + runInfo3.firstHalfQSeqRealSize * constInfo.subBlockIdx; + int64_t kvHeadsOffset = runInfo3.kvHeadsOuterIdx * groupSize * embed; + int64_t embedOffset = runInfo3.groupIdx * embed; + + auto attenOutGmTile = GetTile( + attentionOutGmWithLayout, + tla::MakeCoord( + bOffset + qSeqOffset, + kvHeadsOffset + embedOffset), // batch * qSeqlen, kvHeads* groupSize * embed + tla::MakeShape(runInfo3.halfQSeqRealSize, embedTemplateType)); + epilogueRescaleO( + attenOutGmTile, expUb[taskIdMod3], sumUb[multiCoreIdxMod3], bmm2Tensor, isFirstLoop, + isLastUpdate, MM2_RES_INTRA_EVENT[taskIdMod2]); + } + } + ++taskId; + } + ++multiCoreInnerIdx; + } + qSeqAxisStartIdx = 0; + } + } + +private: + static constexpr uint32_t embedTemplateAlign64 = Align64Func((uint16_t)embedTemplateType); + static constexpr uint32_t MM1_RESULT_SIZE = + qSeqlenTemplateType / CV_RATIO * kvSeqlenTemplateType * sizeof(ElementS); + static constexpr uint32_t MM2_RESULT_SIZE = + qSeqlenTemplateType / CV_RATIO * embedTemplateAlign64 * sizeof(ElementOTmp); + static constexpr uint32_t SHARE_UB_SIZE = CeilDiv(qSeqlenTemplateType, NUM2) * sizeof(ElementS); + + AscendC::LocalTensor bmm1TensorList[NUM2]; + AscendC::LocalTensor mm2AL1TensorList[KERNEL_TASK_NUM]; + AscendC::LocalTensor bmm2TensorList[NUM2]; + AscendC::LocalTensor expUb[KERNEL_TASK_NUM]; + AscendC::LocalTensor sumUb[KERNEL_TASK_NUM]; + AscendC::LocalTensor maxUb[KERNEL_TASK_NUM]; + ConstInfo constInfo; + AttenMaskInfo attenMaskInfo; + uint32_t blockIdx; + uint32_t subBlockIdx; + + RunInfo runInfo[4]; // 最内层循环kvSeq参数 + RunParamStr runParam; // 外层参数 + uint32_t l1BufAddrStart = 0; + uint32_t ubBufAddrStart = 0; + + Arch::Resource resource; + + /* =====================运行时变量==================== */ + CubeCoordInfo coordInfo[3]; + + // =========================================== private functions =========================================== + CATLASS_DEVICE void SetRunInfo( + RunInfo& runInfo, RunParamStr& runParam, int64_t taskId, int64_t kvSeqLoopCount, int64_t kvSeqLoopLimit, + int64_t multiCoreInnerIdx) + { + runInfo.kvSeqAxisStartIdx = runParam.kvSeqAxisLineStartIdx; + runInfo.kvSeqLoopStartIdx = runParam.kvSeqLoopStartIdx; + runInfo.kvSeqAxisEndIdx = runParam.kvSeqAxisLineEndIdx; + runInfo.kvSeqLoopCount = kvSeqLoopCount; + if (runInfo.multiCoreInnerIdx != multiCoreInnerIdx) { + runInfo.qSeqOuterAxisIdx = runParam.qSeqOuterAxisIdx; + runInfo.batchOuterIdx = runParam.batchOuterIdx; + runInfo.kvHeadsOuterIdx = runParam.kvHeadsOuterIdx; + runInfo.groupIdx = runParam.groupIdx; + runInfo.multiCoreInnerIdx = multiCoreInnerIdx; + runInfo.multiCoreIdxMod2 = multiCoreInnerIdx & 1; + runInfo.multiCoreIdxMod3 = multiCoreInnerIdx % 3; + } + + runInfo.taskId = taskId; + runInfo.taskIdMod2 = taskId & 1; + runInfo.taskIdMod3 = taskId % 3; + runInfo.kvSeqLoopLimit = kvSeqLoopLimit; + + runInfo.actualQSeqSize = runParam.actualQSeqSize; + runInfo.actualKvSeqSize = runParam.actualKvSeqSize; + this->ComputeBmm1Tail(runInfo, runParam); + runInfo.batchOuterIdx = runParam.batchOuterIdx; + } + + CATLASS_DEVICE void ComputeBmm1Tail(RunInfo& runInfo, RunParamStr& runParam) + { + // ------------------------qSeq Base Related--------------------------- + runInfo.qSeqRealSize = runParam.qSeqRealSize; + runInfo.halfQSeqRealSize = runParam.halfQSeqRealSize; + runInfo.firstHalfQSeqRealSize = runParam.firstHalfQSeqRealSize; + + // ------------------------kvSeq Base Related---------------------------- + runInfo.kvSeqRealSize = this->constInfo.kvSeqlenBase; + if ((runInfo.kvSeqLoopCount + 1) * runInfo.kvSeqRealSize > runInfo.kvSeqAxisEndIdx) { + runInfo.kvSeqRealSize = runInfo.kvSeqAxisEndIdx - runInfo.kvSeqLoopCount * runInfo.kvSeqRealSize; + } + } + + CATLASS_DEVICE void CalcQSeqCoord(RunInfo& runInfo, ConstInfo& constInfo) + { + // 计算qSeq方向偏移 + coordInfo[runInfo.taskIdMod3].qSeqCoord = runInfo.qSeqOuterAxisIdx * this->constInfo.qSeqlenBase; + } + + CATLASS_DEVICE void CalcKvSeqCoord(RunInfo& runInfo, ConstInfo& constInfo) + { + coordInfo[runInfo.taskIdMod3].kvSeqCoord = + runInfo.kvSeqAxisStartIdx + + (runInfo.kvSeqLoopCount - runInfo.kvSeqLoopStartIdx) * this->constInfo.kvSeqlenBase; + coordInfo[runInfo.taskIdMod3].curBIdx = runInfo.batchOuterIdx; + } +}; + +template +CATLASS_GLOBAL void FAInferTla( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR mask, GM_ADDR blockTables, GM_ADDR o, GM_ADDR actualQSeqlen, + GM_ADDR actualKvSeqlen, GM_ADDR tiling) +{ + using ArchTag = Arch::Ascend950; + using ElementQ = Dtype; + using LayoutTagQ = layout::RowMajor; + using ElementK = Dtype; + using LayoutTagK = layout::ColumnMajor; + using ElementV = Dtype; + using LayoutTagV = layout::RowMajor; + using ElementS = float; + using LayoutTagS = layout::RowMajor; + using ElementP = Dtype; + using LayoutTagP = layout::zN; + using ElementO = Dtype; + using LayoutTagO = layout::RowMajor; + using ElementMask = uint8_t; + using LayoutTagMask = layout::RowMajor; + using ElementOTmp = float; + using LayoutTagOTmp = layout::RowMajor; + // L1TileShape::K must be embdding + using L1TileShape = tla::Shape<_128, _128, _128>; + using L0TileShape = L1TileShape; + // GEMM Block模块，实现Flash Attention Infer的Q * K^T + using DispatchPolicyQK = Gemm::MmadFAIQK; + using TileCopyQK = Gemm::Tile::PackedTileCopyTlaToUB< + ArchTag, ElementQ, LayoutTagQ, ElementK, LayoutTagK, ElementS, LayoutTagS, void, + Gemm::Tile::CopyL0CToUBMode::SPLIT_M>; + using TileMmadQK = Gemm::Tile::TileMmadTla; + using BlockMmadQK = Gemm::Block::BlockMmadTla< + DispatchPolicyQK, L1TileShape, L0TileShape, ElementQ, ElementK, ElementS, void, TileCopyQK, TileMmadQK>; + + // Epilogue Block模块，实现Flash Attention Infer中当前S基块的softmax + using DispatchPolicySoftmax = Epilogue::EpilogueAscend950FASoftmax; + using PType = Gemm::GemmType; + using SType = Gemm::GemmType; + using maskType = Gemm::GemmType; + using EpilogueOnlineSoftmax = + Epilogue::Block::BlockEpilogue; + + // GEMM Block模块，实现Flash Attention Infer的P * V + using DispatchPolicyPV = Gemm::MmadFAIPV; + using TileCopyPV = Gemm::Tile::PackedTileCopyTlaToUB< + ArchTag, ElementP, LayoutTagP, ElementV, LayoutTagV, ElementOTmp, LayoutTagV, void, + Gemm::Tile::CopyL0CToUBMode::SPLIT_M>; + using TileMmadPV = Gemm::Tile::TileMmadTla; + using BlockMmadPV = Gemm::Block::BlockMmadTla< + DispatchPolicyPV, L1TileShape, L0TileShape, ElementP, ElementV, ElementOTmp, void, TileCopyPV, TileMmadPV>; + + // Epilogue Block模块，实现Flash Attention Infer中当前O基块的更新 + using DispatchPolicyRescaleO = Epilogue::EpilogueAscend950FARescaleO; + using OType = Gemm::GemmType; + using OTmpType = Gemm::GemmType; + using EpilogueRescaleO = Epilogue::Block::BlockEpilogue; + + using FAInferKernel = + FAInferKernel; + FAIKernelParams params{q, k, v, mask, blockTables, actualQSeqlen, actualKvSeqlen, o, tiling}; + // call kernel + FAInferKernel flashAttnInfer; + flashAttnInfer(params); +} + +#endif // XLLM_OPS_XFAI_ARCH35_A5_X_FLASH_ATTENTION_INFER_KERNEL_H \ No newline at end of file diff --git a/xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer_kernel_utils.h b/xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer_kernel_utils.h new file mode 100644 index 0000000..4e60c13 --- /dev/null +++ b/xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer_kernel_utils.h @@ -0,0 +1,233 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file a5_x_flash_attention_infer_kernel_utils.h + * \brief + */ + +#ifndef XLLM_OPS_XFAI_ARCH35_A5_X_FLASH_ATTENTION_INFER_KERNEL_UTILS_H +#define XLLM_OPS_XFAI_ARCH35_A5_X_FLASH_ATTENTION_INFER_KERNEL_UTILS_H + +#include "catlass/catlass.hpp" +using namespace Catlass; +using namespace AscendC; + +constexpr uint32_t CV_RATIO = 2; +constexpr uint32_t NUM2 = 2; +constexpr uint32_t KERNEL_TASK_NUM = 3; + +template +CATLASS_DEVICE T Min(T a, T b) +{ + return (a > b) ? b : a; +} + +struct FAIKernelParams { + GM_ADDR q; + GM_ADDR k; + GM_ADDR v; + GM_ADDR mask; + GM_ADDR blockTables; + GM_ADDR actualQSeqlen; + GM_ADDR actualKvSeqlen; + GM_ADDR o; + GM_ADDR tiling; + // Methods + CATLASS_DEVICE + FAIKernelParams() + {} + CATLASS_DEVICE + FAIKernelParams( + GM_ADDR q_, GM_ADDR k_, GM_ADDR v_, GM_ADDR mask_, GM_ADDR blockTables_, GM_ADDR actualQSeqlen_, + GM_ADDR actualKvSeqlen_, GM_ADDR o_, GM_ADDR tiling_) + : q(q_), + k(k_), + v(v_), + mask(mask_), + blockTables(blockTables_), + actualQSeqlen(actualQSeqlen_), + actualKvSeqlen(actualKvSeqlen_), + o(o_), + tiling(tiling_) + {} +}; + +constexpr uint64_t SYNC_MODE = 4; +constexpr uint64_t SYNC_C1_V1_FLAG[2] = {0, 1}; +constexpr uint64_t SYNC_V1_C2_FLAG[3] = {2, 3, 4}; +constexpr uint64_t SYNC_C2_V2_FLAG[2] = {5, 6}; + +constexpr uint64_t MM2_RES_INTRA_EVENT[2] = {7, 8}; // mm2ResIntraEvent +constexpr uint64_t MM1_RES_INTRA_EVENT[2] = {9, 10}; // mm1ResIntraEvent + +struct CubeCoordInfo { + uint32_t curBIdx; + uint32_t qSeqCoord; + uint32_t kvSeqCoord; +}; + +struct RunParamStr { // 分核与切块需要使用到参数 + int64_t batchOuterIdx; + int64_t qSeqOuterAxisIdx; + int64_t kvHeadsOuterIdx; + int64_t groupIdx; + int32_t kvSeqLoopStartIdx; /* kvSeq方向的循环控制信息 souter层确定 */ + int32_t kvSeqLoopEndIdx; /* kvSeq方向的循环控制信息 souter层确定 */ + int64_t kvSeqAxisLineStartIdx = 0; /* kvSeq方向按行的起始位置 */ + int64_t kvSeqAxisLineEndIdx; /* kvSeq方向按行的结束位置 */ + uint32_t qSeqRealSize; + uint32_t halfQSeqRealSize; + uint32_t firstHalfQSeqRealSize; + int64_t actualQSeqSize; /* Q的actualSeqLength */ + int64_t actualKvSeqSize; /* KV的actualSeqLength */ + int64_t qSeqLoopTimes; +}; + +struct RunInfo { + int64_t kvSeqAxisStartIdx; /* kvSeq的起始位置*/ + int64_t kvSeqAxisEndIdx; + int64_t kvSeqLoopCount; /* kvSeq循环当前的循环index */ + int64_t kvSeqLoopStartIdx; + int64_t kvSeqLoopLimit; + int64_t qSeqOuterAxisIdx = 0; /* qSeq轴的index */ + int64_t batchOuterIdx = 0; /* b轴的index */ + int64_t kvHeadsOuterIdx = 0; /* n2轴的index */ + int64_t groupIdx = 0; /* g轴的index */ + int32_t qSeqRealSize; + int32_t halfQSeqRealSize; /* vector侧实际的qSeq基本块大小，如果Cube基本块=128，那么halfQSeqRealSize=64 */ + int32_t + firstHalfQSeqRealSize; /* 当qSeqRealSize不是2的整数倍时，v0比v1少计算一行，计算subblock偏移的时候需要使用v0的qSeq + size */ + int32_t kvSeqRealSize; /* kvSeq方向基本块的真实长度 */ + int64_t taskId; + int64_t multiCoreInnerIdx = 0; + int64_t actualQSeqSize; /* 非TND场景=总qSeqSize, Tnd场景下当前batch对应的qSeq */ + int64_t actualKvSeqSize; /* 非TND场景=总kvSeqSize, Tnd场景下当前batch对应的kvSeq */ + uint8_t taskIdMod2; + uint8_t taskIdMod3; + uint8_t multiCoreIdxMod2 = 0; + uint8_t multiCoreIdxMod3 = 0; + int64_t blockTableOffset; +}; + +struct ConstInfo { + /* 全局的基本块信息 */ + uint32_t qSeqlenBase; + uint32_t kvSeqlenBase; + int64_t embed; + int64_t groupSize; /* g轴的大小 */ + int64_t qHeads; + int64_t kvHeads; + int64_t qSeqlen; /* qSeq总大小 */ + int64_t kvSeqlen; /* kvSeq总大小 */ + /* 轴的乘积 */ + int64_t qSeqlenOuterSize; + uint8_t subBlockIdx; + float scaleValue; + /* 推理新增 */ + bool isActualLenDimsNull; /* 判断是否有actualseq */ + bool isActualLenDimsKVNull; /* 判断是否有actualseq_kv */ + uint32_t actualSeqLenSize; /* 用户输入的actualseq的长度 */ + uint32_t actualSeqLenKVSize; /* 用户输入的actualseq_kv的长度 */ + /* service mm1 mm2 pageAttention */ + uint32_t blockTableDim2; + uint32_t blockSize; + uint32_t paBlockNumSum; + /* G S不合轴场景，外层循环是B、N2、G，内层循环S，headNumRatio = groupSize */ + uint32_t headNumRatio; + uint32_t bnAxisStartIdx; + uint32_t bnAxisEndIdx; + uint32_t actualSeqLengthsSize; + uint32_t actualSeqLengthsKVSize; + bool isActualSeqLengthsNull; + bool isActualSeqLengthsKVNull; + /* base params */ + uint32_t batch; + /* special params */ + uint32_t attenMaskQSeqlen; + uint32_t attenMaskKvSeqlen; + /* core params */ + volatile int64_t multiCoreInnerOffset; /* 二次赋值的变量需要volatile修饰 */ + volatile int64_t multiCoreInnerLimit; /* 二次赋值的变量需要volatile修饰 */ + uint32_t coreNum; +}; + +struct AttenMaskInfo { + int64_t attenMaskShapeType; + int64_t attenMaskQSeqlen; + int64_t attenMaskKvSeqlen; + int64_t attenMaskOffsetPre; +}; + +constexpr uint16_t SHIFT_NUM_6 = 6; +constexpr uint16_t ADD_NUM_63 = 63; +CATLASS_DEVICE constexpr uint16_t Align64Func(uint16_t data) +{ + return (data + ADD_NUM_63) >> SHIFT_NUM_6 << SHIFT_NUM_6; +} +CATLASS_DEVICE constexpr uint16_t Align(uint16_t data, uint16_t baseSize) +{ + return (data - 1) / baseSize * baseSize + baseSize; +} + +CATLASS_DEVICE void ComputeParamBatch( + RunParamStr& runParam, const ConstInfo& constInfo, const AttenMaskInfo& attenMaskInfo) +{ + runParam.actualQSeqSize = constInfo.qSeqlen; + ; + runParam.actualKvSeqSize = constInfo.kvSeqlen; + ; +} + +template +CATLASS_DEVICE void ComputeQseqLoopInfo( + RunParamStr& runParam, const ConstInfo& constInfo, bool lastBN, int64_t nextQSeqAxisIdx) +{ + constexpr int32_t qSeqlenBase = static_cast(qSeqlenTemplateType); + int32_t qSeqLoopTimes = CeilDiv(runParam.actualQSeqSize, qSeqlenBase); + // 不是最后一个bn, 赋值souterBlockNum + if (!lastBN) { + runParam.qSeqLoopTimes = qSeqLoopTimes; + } else { // 最后一个bn, 从数组下一个元素取值 + runParam.qSeqLoopTimes = nextQSeqAxisIdx == 0 ? qSeqLoopTimes : nextQSeqAxisIdx; + } +} + +template +CATLASS_DEVICE void ComputeParamQSeq(RunParamStr& runParam, const ConstInfo& constInfo, uint32_t sOuterLoopIdx) +{ + int64_t cubeSOuterOffset = sOuterLoopIdx * (uint32_t)qSeqlenTemplateType; + if (runParam.actualQSeqSize == 0) { + runParam.qSeqRealSize = 0; + } else { + runParam.qSeqRealSize = + Min((uint32_t)qSeqlenTemplateType, (uint32_t)(runParam.actualQSeqSize - cubeSOuterOffset)); + } + + runParam.halfQSeqRealSize = (runParam.qSeqRealSize + 1) >> 1; + runParam.firstHalfQSeqRealSize = runParam.halfQSeqRealSize; + if (constInfo.subBlockIdx == 1) { + runParam.halfQSeqRealSize = runParam.qSeqRealSize - runParam.halfQSeqRealSize; + } +} + +template +CATLASS_DEVICE void ComputeKvSeqLoopInfo(RunParamStr& runParam, const ConstInfo& constInfo) +{ + constexpr int32_t kvSeqlenBase = static_cast(kvSeqlenTemplateType); + runParam.kvSeqAxisLineStartIdx = 0; + runParam.kvSeqAxisLineEndIdx = runParam.actualKvSeqSize; + runParam.kvSeqLoopStartIdx = 0; + runParam.kvSeqLoopEndIdx = (runParam.kvSeqAxisLineEndIdx + kvSeqlenBase - 1) / kvSeqlenBase; +} + +#endif // XLLM_OPS_XFAI_ARCH35_A5_X_FLASH_ATTENTION_INFER_KERNEL_UTILS_H \ No newline at end of file diff --git a/xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer_tiling_data_def.h b/xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer_tiling_data_def.h new file mode 100644 index 0000000..df0aa4f --- /dev/null +++ b/xllm_ops/x_flash_attention_infer/op_kernel/arch35/a5_x_flash_attention_infer_tiling_data_def.h @@ -0,0 +1,59 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-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 XLLM_OPS_XFAI_ARCH35_A5_X_FLASH_ATTENTION_INFER_TILING_DATA_DEF_H +#define XLLM_OPS_XFAI_ARCH35_A5_X_FLASH_ATTENTION_INFER_TILING_DATA_DEF_H + +constexpr uint32_t MAX_CORE_NUM = 64; + +class InputParamsRegbase { +public: + int64_t batch; + int64_t qHeads; + int64_t kvHeads; + int64_t groupSize; + int64_t qSeqlen; + int64_t kvSeqlen; + int64_t embed; + float scaleValue; + uint8_t attenMaskCompressMode; // SPARSE_MODE_NO_MASK: 0, SPARSE_MODE_LEFT_UP: 1, SPARSE_MODE_RIGHT_DOWN : 2 + + // PFA + uint8_t isActualSeqLengthsNull; + uint8_t isActualSeqLengthsKVNull; + uint32_t actualSeqLengthsSize; + uint32_t actualSeqLengthsKVSize; + + uint32_t headNumRatio; + uint32_t blockSize; + uint32_t blockTableDim2; + uint32_t paBlockNumSum; + uint32_t attenMaskQSeqlen; + uint32_t attenMaskKvSeqlen; +}; + +class MultiCoreParamsRegbase { +public: + int32_t coreNum; + int64_t totalSize; + int64_t qSeqlenOuterSize; + int64_t splitFactorSize; + int64_t splitFactorTailSize; + uint32_t bnAxisStartIdx[MAX_CORE_NUM]; + int64_t sparseStartIdx[MAX_CORE_NUM]; +}; + +class FATilingData { +public: + InputParamsRegbase inputParamsRegbase; + MultiCoreParamsRegbase multiCoreParamsRegbase; +}; +#endif // XLLM_OPS_XFAI_ARCH35_A5_X_FLASH_ATTENTION_INFER_TILING_DATA_DEF_H \ No newline at end of file diff --git a/xllm_ops/x_flash_attention_infer/op_kernel/x_flash_attention_infer.cpp b/xllm_ops/x_flash_attention_infer/op_kernel/x_flash_attention_infer.cpp index 107b5cb..82bdce4 100644 --- a/xllm_ops/x_flash_attention_infer/op_kernel/x_flash_attention_infer.cpp +++ b/xllm_ops/x_flash_attention_infer/op_kernel/x_flash_attention_infer.cpp @@ -8,12 +8,75 @@  * See LICENSE in the root of the software repository for the full text of the License.  */ +// A5(Ascend950/DAV_3510) arch guard. +// Device side must use __NPU_ARCH__ (per catlass migration guide); host side uses +// CATLASS_ARCH. Accept either so the A5 path is selected regardless of which macro +// the toolchain injects for the kernel translation unit. +#if (defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)) || (defined(CATLASS_ARCH) && (CATLASS_ARCH == 3510)) +#define XFA_ARCH35 1 +#endif + +// Device(kernel) side lacks -DCATLASS_ARCH (host-only inject). Derive it from +// __NPU_ARCH__ HERE, before ANY include, so every catlass forwarding header in +// this translation unit (incl. common.h below) dispatches to the ascend950 +// specialization consistently. +#if defined(XFA_ARCH35) && !defined(CATLASS_ARCH) +#define CATLASS_ARCH 3510 +#endif + +// A3(AtlasA2/A3, __NPU_ARCH__ == 2201) arch guard. +// The A3 path (non-XFA_ARCH35) still pulls catlass forwarding headers via +// x_flash_attention_infer.h; the new catlass tile-copy forwarders dispatch +// ONLY when CATLASS_ARCH is explicitly 2201/3510 (host-only inject on device). +// Derive CATLASS_ARCH from __NPU_ARCH__ HERE (before ANY include) so the A3 +// device translation unit resolves CopyGmToL1/CopyL1ToL0A/ScaleGranularity/... +// without affecting the already-validated A5(3510) path above. +#if !defined(XFA_ARCH35) && !defined(CATLASS_ARCH) && defined(__NPU_ARCH__) && (__NPU_ARCH__ == 2201) +#define CATLASS_ARCH 2201 +#endif + +#if defined(XFA_ARCH35) +// arch35 (example49-ported) defines its own FAIKernelParams / helpers, which +// clash with x_flash_attention_infer_common.h. So DO NOT pull common.h here; +// instead declare only the TILING_KEY constants this branch dispatches on. +// Values mirror x_flash_attention_infer_common.h:85-88 (keep in sync). +#ifndef QFP16_KVFP16_TND_CAUSALMASK_FD_TILING +#define QFP16_KVFP16_TND_CAUSALMASK_FD_TILING 1000000000000001113 +#endif +#ifndef QFP16_KVFP16_KVNZ_CAUSALMASK_FD_TILING +#define QFP16_KVFP16_KVNZ_CAUSALMASK_FD_TILING 1000000000000001213 +#endif +#ifndef QBF16_KVBF16_TND_CAUSALMASK_FD_TILING +#define QBF16_KVBF16_TND_CAUSALMASK_FD_TILING 1000000000000001123 +#endif +#ifndef QBF16_KVBF16_KVNZ_CAUSALMASK_FD_TILING +#define QBF16_KVBF16_KVNZ_CAUSALMASK_FD_TILING 1000000000000001223 +#endif +#include "arch35/a5_x_flash_attention_infer.h" +#else #include "x_flash_attention_infer.h" #include "x_flash_attention_infer_fd.h" +#endif extern "C" __global__ __aicore__ void x_flash_attention_infer(GM_ADDR query, GM_ADDR key_cache, GM_ADDR value_cache, GM_ADDR mask, GM_ADDR block_table, GM_ADDR actual_q_lens, GM_ADDR actual_kv_lens, GM_ADDR extra_tiling, GM_ADDR attn_out, GM_ADDR workspace, GM_ADDR tiling) { +#if defined(XFA_ARCH35) + // A5(Ascend950/DAV_3510): dispatch into arch35 example49-ported FAInferKernel. + // Host tiling A5 branch (bnAxisStartIdx/sparseStartIdx) is filled in stage-3. + SetAtomicNone(); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + GET_TILING_DATA(tiling_data, tiling); + if (TILING_KEY_IS(QFP16_KVFP16_TND_CAUSALMASK_FD_TILING) || + TILING_KEY_IS(QFP16_KVFP16_KVNZ_CAUSALMASK_FD_TILING)) { + XllmOps::XfaArch35::FAInferA5Dispatch( + query, key_cache, value_cache, mask, block_table, actual_q_lens, actual_kv_lens, attn_out, tiling); + } else if (TILING_KEY_IS(QBF16_KVBF16_TND_CAUSALMASK_FD_TILING) || + TILING_KEY_IS(QBF16_KVBF16_KVNZ_CAUSALMASK_FD_TILING)) { + XllmOps::XfaArch35::FAInferA5Dispatch( + query, key_cache, value_cache, mask, block_table, actual_q_lens, actual_kv_lens, attn_out, tiling); + } +#else // workspace use; [s,p,oTemp,oUpdate,shared_workspace,unshared_workspace] SetAtomicNone(); KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); @@ -56,4 +119,5 @@ extern "C" __global__ __aicore__ void x_flash_attention_infer(GM_ADDR query, GM_ AscendC::SyncAll(); } } +#endif } From 1a7b356f277ab10caa418d88cedd179ef6d23e0e Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Fri, 31 Jul 2026 16:01:15 +0800 Subject: [PATCH 02/24] feat: bugfix --- test/cpp_test/CMakeLists.txt | 75 ++++++++++--------- test/cpp_test/group_gemm.h | 46 +++++++++++- test/cpp_test/utils_tensor.h | 12 +++ xllm_ops/build_aclnn.sh | 2 +- .../op_kernel/x_attention_catlass_kernel.h | 14 ++++ .../op_kernel/x_attention_tl.cpp | 13 ++++ 6 files changed, 123 insertions(+), 39 deletions(-) diff --git a/test/cpp_test/CMakeLists.txt b/test/cpp_test/CMakeLists.txt index 0e8ecc7..a1c7d46 100644 --- a/test/cpp_test/CMakeLists.txt +++ b/test/cpp_test/CMakeLists.txt @@ -83,21 +83,34 @@ else() endif() set(INCLUDE_BASE_DIR "${ASCEND_PATH}/include") -set(OP_API_PATH "/usr/local/Ascend/ascend-toolkit/latest/opp/vendors/xllm/op_api") +set(OP_API_PATH "${ASCEND_PATH}/opp/vendors/custom_xllm_math/op_api") + +# Platform detection: A5(Ascend950) CANN provides the dedicated NZ-weight +# grouped matmul interface header, while A3(Ascend910_93) does not. +# Use its presence to pick the correct native reference (golden) interface. +if(EXISTS "${INCLUDE_BASE_DIR}/aclnnop/aclnn_grouped_matmul_weight_nz.h") + add_compile_definitions(USE_GROUPED_MATMUL_WEIGHT_NZ) + message(STATUS "Detected aclnnGroupedMatmulWeightNz header: golden path uses WeightNz interface (A5)") +else() + message(STATUS "aclnnGroupedMatmulWeightNz header not found: golden path uses GroupedMatmulV4 interface (A3)") +endif() +set(PYTHON_DIR "python$ENV{PYTHON_VERSION}") +set(PYTHON_LIB_DIR "/usr/local/${PYTHON_DIR}/lib/python3.11/site-packages") # Common include directories/compile options/link options/libraries set(COMMON_INCLUDE_DIRS - "${INCLUDE_BASE_DIR}" - "${INCLUDE_BASE_DIR}/aclnn" - "${INCLUDE_BASE_DIR}/aclnn/op_dev" - "${INCLUDE_BASE_DIR}/platform" - "${ASCEND_PATH}/x86_64-linux/include/exe_graph/runtime" - "/usr/local/lib64/python3.11/site-packages/torch_npu/include" - "/usr/local/lib64/python3.11/site-packages/torch/include" - "/usr/local/lib64/python3.11/site-packages/torch/include/torch/csrc/api/include" - "/usr/local/libtorch_npu/include" - "${OP_API_PATH}/include" - "${CMAKE_CURRENT_SOURCE_DIR}" + "${INCLUDE_BASE_DIR}" + "${INCLUDE_BASE_DIR}/aclnn" + "${INCLUDE_BASE_DIR}/aclnn/op_dev" + "${INCLUDE_BASE_DIR}/platform" + "${ASCEND_PATH}/$ENV{VCPKG_TARGET_TRIPLET}/include/exe_graph/runtime" + "${PYTHON_LIB_DIR}/torch/include/torch/csrc/api/include" + "${PYTHON_LIB_DIR}/torch_npu/include" + "${PYTHON_LIB_DIR}/torch/include" + "${PYTHON_LIB_DIR}/torch/include/torch/csrc/api/include" + "/usr/local/libtorch_npu/include" + "${OP_API_PATH}/include/aclnnop" + "${CMAKE_CURRENT_SOURCE_DIR}" ) @@ -113,19 +126,20 @@ set(COMMON_LINK_OPTS ) set(COMMON_LIBS - "${ASCEND_PATH}/lib64/libascendcl.so" - "${ASCEND_PATH}/lib64/libnnopbase.so" - "${ASCEND_PATH}/lib64/libacl_op_compiler.so" - "${ASCEND_PATH}/lib64/libascendalog.so" - "${ASCEND_PATH}/lib64/libtiling_api.a" - "${ASCEND_PATH}/lib64/libplatform.so" - "${OP_API_PATH}/lib/libcust_opapi.so" - "${ASCEND_PATH}/lib64/libopapi.so" - "/usr/local/lib64/python3.11/site-packages/torch/lib/libc10.so" - "/usr/local/lib64/python3.11/site-packages/torch/lib/libtorch.so" - "/usr/local/lib64/python3.11/site-packages/torch/lib/libtorch_cpu.so" - "/usr/local/libtorch_npu/lib/libtorch_npu.so" - stdc++ + "${ASCEND_PATH}/lib64/libascendcl.so" + "${ASCEND_PATH}/lib64/libnnopbase.so" + "${ASCEND_PATH}/lib64/libacl_op_compiler.so" + "${ASCEND_PATH}/lib64/libascendalog.so" + "${ASCEND_PATH}/lib64/libtiling_api.a" + "${ASCEND_PATH}/lib64/libplatform.so" + "${OP_API_PATH}/lib/libcust_opapi.so" + "${ASCEND_PATH}/lib64/libopapi.so" + "${PYTHON_LIB_DIR}/torch/lib/libc10.so" + "${PYTHON_LIB_DIR}/torch/lib/libtorch.so" + "${PYTHON_LIB_DIR}/torch/lib/libtorch_cpu.so" + "/usr/local/libtorch_npu/lib/libtorch_npu.so" + + stdc++ ) # ============================================================================= @@ -154,15 +168,6 @@ if(ENABLE_PCH) endif() # GTest version tests -add_executable(pp_matmul_test - pp_matmul_test.cpp -) -target_link_libraries(pp_matmul_test PRIVATE - aclnn_common - GTest::gtest - GTest::gtest_main -) - # Group GEMM GTest version add_executable(group_gemm_gtest group_gemm_test.cpp @@ -205,7 +210,6 @@ target_link_libraries(convert_kv_cache_format_test PRIVATE # Add tests # add_test(AllTestsInBeamSearch beam_search_test) -add_test(AllTestsInPPMatmul pp_matmul_test) add_test(AllTestsInGroupGemm group_gemm_gtest) add_test(AllTestsInMultiLatentAttention multi_latent_attention_gtest) add_test(AllTestsInConvertKvCacheFormat convert_kv_cache_format_test) @@ -213,7 +217,6 @@ add_test(AllTestsInConvertKvCacheFormat convert_kv_cache_format_test) # GoogleTest automatic discovery include(GoogleTest) # gtest_discover_tests(beam_search_test) -gtest_discover_tests(pp_matmul_test) gtest_discover_tests(group_gemm_gtest) gtest_discover_tests(multi_latent_attention_gtest) gtest_discover_tests(convert_kv_cache_format_test) diff --git a/test/cpp_test/group_gemm.h b/test/cpp_test/group_gemm.h index d188f0d..0ce3a9c 100644 --- a/test/cpp_test/group_gemm.h +++ b/test/cpp_test/group_gemm.h @@ -18,7 +18,11 @@ limitations under the License. #define GROUP_GEMM_H #include "aclnn_index_group_matmul.h" +#ifdef USE_GROUPED_MATMUL_WEIGHT_NZ +#include "aclnnop/aclnn_grouped_matmul_weight_nz.h" +#else #include "aclnnop/aclnn_grouped_matmul_v4.h" +#endif #include "utils_print.h" #include "utils_tensor.h" namespace group_gemm { @@ -331,6 +335,36 @@ class GroupGemmNative { int64_t groupListType = 0; int64_t actType = 0; +#ifdef USE_GROUPED_MATMUL_WEIGHT_NZ + // A5(Ascend950): NZ-format weight must go through the dedicated + // aclnnGroupedMatmulWeightNz interface. Compared with V4 it has two extra + // params after actType: tuningConfigOptional(nullptr) and quantGroupSize(0). + aclIntArray* tuningConfig = nullptr; + int64_t quantGroupSize = 0; + auto ret = aclnnGroupedMatmulWeightNzGetWorkspaceSize(x, + weight, + bias, + scale, + offset, + antiquantScale, + antiquantOffset, + perTokenScale, + groupedList, + activationInput, + activationQuantScale, + activationQuantOffset, + splitItem, + groupType, + groupListType, + actType, + tuningConfig, + quantGroupSize, + y, + activationFeatureOut, + dynQuantScaleOut, + &workspaceSize, + &executor); +#else auto ret = aclnnGroupedMatmulV4GetWorkspaceSize(x, weight, bias, @@ -351,7 +385,8 @@ class GroupGemmNative { activationFeatureOut, dynQuantScaleOut, &workspaceSize, - &executor); + &executor); +#endif CHECK_RET( ret == ACL_SUCCESS, LOG_PRINT("aclnnGroupedMatmulGetWorkspaceSize failed. ERROR: %d\n", @@ -367,10 +402,17 @@ class GroupGemmNative { return ret); } +#ifdef USE_GROUPED_MATMUL_WEIGHT_NZ + ret = aclnnGroupedMatmulWeightNz(workspaceAddr, workspaceSize, executor, stream); + CHECK_RET(ret == ACL_SUCCESS, + LOG_PRINT("aclnnGroupedMatmulWeightNz failed. ERROR: %d\n", ret); + return ret); +#else ret = aclnnGroupedMatmulV4(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, - LOG_PRINT("aclnnIndexGroupMatmul failed. ERROR: %d\n", ret); + LOG_PRINT("aclnnGroupedMatmulV4 failed. ERROR: %d\n", ret); return ret); +#endif ret = aclrtSynchronizeStream(stream); CHECK_RET(ret == ACL_SUCCESS, diff --git a/test/cpp_test/utils_tensor.h b/test/cpp_test/utils_tensor.h index edc25ae..ad137d4 100644 --- a/test/cpp_test/utils_tensor.h +++ b/test/cpp_test/utils_tensor.h @@ -93,11 +93,23 @@ int Init(int32_t deviceId, aclrtStream* stream) { std::vector get_weight_storage_shape(const std::vector& shape) { std::vector storageTensorDims (5, 0); // ND格式下，storageShape和originalShape一致 +#ifdef USE_GROUPED_MATMUL_WEIGHT_NZ + // A5(Ascend950/DAV_3510): INT8 weight 的 FRACTAL_NZ 分形内轴为 16x32 + // 校验器要求 storage shape = [g, ceil(n/32), ceil(k/16), 16, 32] + // 此处 trans_shape = {g, k, n}，故 shape[1]=k, shape[2]=n + storageTensorDims[0] = shape[0]; + storageTensorDims[1] = 1 + ((shape[2] - 1) / 32); // ceil(n/32)：INT8 NZ 外轴 n + storageTensorDims[2] = 1 + ((shape[1] - 1) / 16); // ceil(k/16)：内轴 k + storageTensorDims[3] = 16; // 3, 16：NZ格式要求 + storageTensorDims[4] = 32; // 4, 32：INT8 NZ格式内轴要求 +#else + // A2/A3: FP16/INT8 均使用 16x16 分形 storageTensorDims[0] = shape[0]; storageTensorDims[1] = 1 + ((shape[1] - 1) / 16); // 1, 16：1: 维度, 16: padding大小 storageTensorDims[2] = 1 + ((shape[2] - 1) / 16); // 2, 16：1: 维度, 16: padding大小 storageTensorDims[3] = 16; // 3, 16：NZ格式要求 storageTensorDims[4] = 16; // 4, 16：NZ格式要求 +#endif return storageTensorDims; } diff --git a/xllm_ops/build_aclnn.sh b/xllm_ops/build_aclnn.sh index 9851ec6..e62f2fc 100644 --- a/xllm_ops/build_aclnn.sh +++ b/xllm_ops/build_aclnn.sh @@ -323,7 +323,7 @@ elif [[ "$SOC_VERSION" =~ ^ascend950 ]]; then "hc_post" "rms_norm_dynamic_quant" "inplace_partial_rotary_mul" - "dispatch_ffn_combine" + #"dispatch_ffn_combine" "dequant_swiglu_quant" ## 已在 CANN 中内置，删除后会有精度问题，CANN内置见 aarch64-linux/include/aclnnop/aclnn_dequant_swiglu_quant.h "scatter_nd_update_v2" diff --git a/xllm_ops/x_attention/op_kernel/x_attention_catlass_kernel.h b/xllm_ops/x_attention/op_kernel/x_attention_catlass_kernel.h index 78bb7ec..6add9bc 100644 --- a/xllm_ops/x_attention/op_kernel/x_attention_catlass_kernel.h +++ b/xllm_ops/x_attention/op_kernel/x_attention_catlass_kernel.h @@ -16,6 +16,20 @@ limitations under the License. #ifndef X_ATTN_CATLASS_KERNEL_H #define X_ATTN_CATLASS_KERNEL_H +// [catlass arch guard] +// The catlass tile-copy forwarding headers (e.g. gemm/tile/copy_gm_to_l1.hpp) +// dispatch ONLY when CATLASS_ARCH is explicitly defined as 2201 (AtlasA2/A3) or +// 3510 (Ascend950/A5); otherwise CopyGmToL1/CopyL1ToL0A/... templates are never +// defined. The kernel(device) translation unit is NOT given -DCATLASS_ARCH +// (host-only inject), but the toolchain injects __NPU_ARCH__ (2201 for A2/A3, +// 3510 for A5). Derive CATLASS_ARCH from it HERE, before ANY catlass include, +// so this fix works on A3 without affecting the already-validated A5 path. +#if !defined(CATLASS_ARCH) && defined(__NPU_ARCH__) +#if (__NPU_ARCH__ == 2201) || (__NPU_ARCH__ == 3510) +#define CATLASS_ARCH __NPU_ARCH__ +#endif +#endif + #include "catlass/arch/arch.hpp" #include "catlass/arch/cross_core_sync.hpp" #include "catlass/arch/resource.hpp" diff --git a/xllm_ops/x_attention_tl/op_kernel/x_attention_tl.cpp b/xllm_ops/x_attention_tl/op_kernel/x_attention_tl.cpp index f857dd1..e59105b 100644 --- a/xllm_ops/x_attention_tl/op_kernel/x_attention_tl.cpp +++ b/xllm_ops/x_attention_tl/op_kernel/x_attention_tl.cpp @@ -13,6 +13,19 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +// [catlass arch guard] +// x_attention_tl.h pulls catlass forwarding headers. The new catlass tile-copy +// forwarders dispatch ONLY when CATLASS_ARCH is explicitly 2201 (AtlasA2/A3) or +// 3510 (Ascend950/A5); the device translation unit is NOT given -DCATLASS_ARCH +// (host-only inject) but the toolchain injects __NPU_ARCH__. Derive CATLASS_ARCH +// from it HERE, before ANY include, so A3 resolves CopyGmToL1/CopyL1ToL0A/... +// without affecting the already-validated A5(3510) path. +#if !defined(CATLASS_ARCH) && defined(__NPU_ARCH__) +#if (__NPU_ARCH__ == 2201) || (__NPU_ARCH__ == 3510) +#define CATLASS_ARCH __NPU_ARCH__ +#endif +#endif + #include "acl/acl.h" #include "kernel_operator.h" #include "lib/matmul_intf.h" From 3d9ca8fb14aae2c1ee3cfdecba7bd7ee4a6bfbd1 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Fri, 31 Jul 2026 16:26:48 +0800 Subject: [PATCH 03/24] feat: change the include dir --- test/cpp_test/CMakeLists.txt | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/test/cpp_test/CMakeLists.txt b/test/cpp_test/CMakeLists.txt index a1c7d46..567cc31 100644 --- a/test/cpp_test/CMakeLists.txt +++ b/test/cpp_test/CMakeLists.txt @@ -96,6 +96,10 @@ else() endif() set(PYTHON_DIR "python$ENV{PYTHON_VERSION}") set(PYTHON_LIB_DIR "/usr/local/${PYTHON_DIR}/lib/python3.11/site-packages") +set(ASCEND_CANN_INCLUDE_NAME $ENV{VCPKG_TARGET_TRIPLET}) +if("$ENV{VCPKG_TARGET_TRIPLET}" MATCHES "x64-linux") + set(ASCEND_CANN_INCLUDE_NAME "x86_64-linux") +endif() # Common include directories/compile options/link options/libraries set(COMMON_INCLUDE_DIRS @@ -103,7 +107,7 @@ set(COMMON_INCLUDE_DIRS "${INCLUDE_BASE_DIR}/aclnn" "${INCLUDE_BASE_DIR}/aclnn/op_dev" "${INCLUDE_BASE_DIR}/platform" - "${ASCEND_PATH}/$ENV{VCPKG_TARGET_TRIPLET}/include/exe_graph/runtime" + "${ASCEND_PATH}/${ASCEND_CANN_INCLUDE_NAME}/include/exe_graph/runtime" "${PYTHON_LIB_DIR}/torch/include/torch/csrc/api/include" "${PYTHON_LIB_DIR}/torch_npu/include" "${PYTHON_LIB_DIR}/torch/include" From 2d2a361a5f6c71a280d6d1b4c637dfae5f1ef04e Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Fri, 31 Jul 2026 16:29:21 +0800 Subject: [PATCH 04/24] feat: adapte the a5 --- .../op_host/multi_latent_attention_tiling.cpp | 7 ------- 1 file changed, 7 deletions(-) diff --git a/xllm_ops/multi_latent_attention/op_host/multi_latent_attention_tiling.cpp b/xllm_ops/multi_latent_attention/op_host/multi_latent_attention_tiling.cpp index 5993ac0..19c9fe0 100644 --- a/xllm_ops/multi_latent_attention/op_host/multi_latent_attention_tiling.cpp +++ b/xllm_ops/multi_latent_attention/op_host/multi_latent_attention_tiling.cpp @@ -18,16 +18,9 @@ limitations under the License. #include "tiling/platform/platform_ascendc.h" namespace optiling { - #ifdef OP_TILING_LIB static ge::graphStatus TilingFunc(gert::TilingContext *context) { return AtbOps::MLATiling(context); - // return ge::GRAPH_SUCCESS; } - #else - static ge::graphStatus TilingFunc(gert::TilingContext *context) { - return ge::GRAPH_SUCCESS; - } - #endif IMPL_OP_OPTILING(MultiLatentAttention) .Tiling(TilingFunc); } From bd471dd7c0638986c459698265b5d170fdd23fec Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Fri, 31 Jul 2026 16:40:11 +0800 Subject: [PATCH 05/24] feat: slove the multi_latent_attention --- .../test_multi_latent_attention.py | 23 +++++++++++++++---- .../multi_latent_attention_tiling_impl.cpp | 8 +++---- 2 files changed, 22 insertions(+), 9 deletions(-) diff --git a/test/python_test/test_multi_latent_attention.py b/test/python_test/test_multi_latent_attention.py index f5514f8..c2f2180 100644 --- a/test/python_test/test_multi_latent_attention.py +++ b/test/python_test/test_multi_latent_attention.py @@ -113,9 +113,22 @@ def _mla_decode_golden(q_nope, q_rope, k_nope, k_rope, v_nope, @pytest.mark.parametrize( "dtype, batch, q_head, kv_head, kv_seqlen, block_size", [ - (torch.float16, 1, 16, 1, 128, 128), - (torch.float16, 2, 16, 1, 128, 128), - (torch.float16, 2, 32, 1, 256, 128), + (torch.float16, 200, 32, 1, 64, 128), + (torch.float16, 1, 128, 1, 1024, 128), + (torch.float16, 6, 128, 1, 2048, 128), + (torch.float16, 12, 128, 1, 2048, 128), + (torch.float16, 24, 128, 1, 4096, 128), + (torch.float16, 25, 128, 1, 4096, 128), + (torch.float16, 1, 32, 1, 1024, 128), + (torch.float16, 6, 32, 1, 2048, 128), + (torch.float16, 12, 32, 1, 2048, 128), + (torch.float16, 24, 32, 1, 4096, 128), + (torch.float16, 25, 32, 1, 4096, 128), + (torch.float16, 1, 64, 1, 1024, 128), + (torch.float16, 6, 64, 1, 2048, 128), + (torch.float16, 12, 64, 1, 2048, 128), + (torch.float16, 24, 64, 1, 4096, 128), + (torch.float16, 25, 64, 1, 4096, 128), ], ) def test_multi_latent_attention(dtype, batch, q_head, kv_head, kv_seqlen, block_size): @@ -151,8 +164,8 @@ def test_multi_latent_attention(dtype, batch, q_head, kv_head, kv_seqlen, block_ out = custom_ops.multi_latent_attention_npu( query.npu(), query_rope.npu(), kv_cache.npu(), kv_cache_rope.npu(), block_table.npu(), context_lens.npu(), - q_head, kv_head, tor, [kv_seqlen] * batch, + q_head, kv_head, tor, [kv_seqlen] * batch, [q_seqlen] * batch, ) out = out.cpu().view(batch, q_head, NOPE_DIM).to(torch.float32) - torch.testing.assert_close(out, golden, atol=6e-2, rtol=6e-2) \ No newline at end of file + torch.testing.assert_close(out, golden, atol=6e-2, rtol=6e-2) diff --git a/xllm_ops/multi_latent_attention/op_host/multi_latent_attention_tiling_impl.cpp b/xllm_ops/multi_latent_attention/op_host/multi_latent_attention_tiling_impl.cpp index 46d023d..59d51e6 100644 --- a/xllm_ops/multi_latent_attention/op_host/multi_latent_attention_tiling_impl.cpp +++ b/xllm_ops/multi_latent_attention/op_host/multi_latent_attention_tiling_impl.cpp @@ -123,18 +123,18 @@ OpParam::MLA GetParamFromTilingContext(gert::TilingContext *context) { auto qSeqLen = context->GetAttrs()->GetListInt(5)->GetData(); size_t arraySize = context->GetAttrs()->GetListInt(5)->GetSize(); param.qSeqLen.reserve(arraySize); - if (arraySize >= 1 && reinterpret_cast(qSeqLen)[0] >= 0) { + if (arraySize >= 1 && reinterpret_cast(qSeqLen)[0] >= 0) { for (size_t i = 0; i < arraySize; ++i) { - param.qSeqLen.push_back(reinterpret_cast(qSeqLen)[i]); + param.qSeqLen.push_back(reinterpret_cast(qSeqLen)[i]); } } auto kvSeqLenAttr = context->GetAttrs()->GetListInt(6)->GetData(); arraySize = context->GetAttrs()->GetListInt(6)->GetSize(); param.kvSeqLen.reserve(arraySize); - if (arraySize >= 1 && reinterpret_cast(kvSeqLenAttr)[0] >= 0) { + if (arraySize >= 1 && reinterpret_cast(kvSeqLenAttr)[0] >= 0) { for (size_t i = 0; i < arraySize; ++i) { - param.kvSeqLen.push_back(reinterpret_cast(kvSeqLenAttr)[i]); + param.kvSeqLen.push_back(reinterpret_cast(kvSeqLenAttr)[i]); } } return param; From 249592ccdfb7af5f7adec851ad11548d3b204bc5 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Wed, 5 Aug 2026 11:19:40 +0800 Subject: [PATCH 06/24] feat: add mult_latent_attention desc --- .../multi_latent_attention_bf16_desc.md | 367 ++++++++++++ .../multi_latent_attention_desc.md | 524 ++++++++++++++++++ 2 files changed, 891 insertions(+) create mode 100644 xllm_ops/multi_latent_attention/multi_latent_attention_bf16_desc.md create mode 100644 xllm_ops/multi_latent_attention/multi_latent_attention_desc.md diff --git a/xllm_ops/multi_latent_attention/multi_latent_attention_bf16_desc.md b/xllm_ops/multi_latent_attention/multi_latent_attention_bf16_desc.md new file mode 100644 index 0000000..0764c77 --- /dev/null +++ b/xllm_ops/multi_latent_attention/multi_latent_attention_bf16_desc.md @@ -0,0 +1,367 @@ +# Multi-Latent Attention (MLA) 算子实现分析(BF16 数据类型） + +> 本文档分析 `xllm_ops/multi_latent_attention` 算子在 Ascend AscendC 平台上、**数据类型为 BF16** 时的实现。 +> - **Host 侧**:聚焦 Tiling 切分策略(tiling 参数、核数计算、任务数、多核分配),与 INT8 基本一致,差异点单独标注。 +> - **Kernel 侧**:聚焦 **BF16 数据类型**的实现(业务处理流程、函数逻辑、数据获取、地址计算),重点对比与 INT8 的差异。 +> - 参考文档:`multi_latent_attention_desc.md`(INT8 版本)。 + +--- + +## 1. 算子概述 + +Multi-Latent Attention(MLA)是 DeepSeek 系列模型使用的注意力机制,核心特点是把 KV Cache 压缩到一个低秩的隐空间(latent),从而大幅降低 KV Cache 显存占用。本算子实现的是 **decode(增量推理)阶段**的 MLA,基于 **PagedAttention** 的 block_table 机制管理 KV Cache。 + +算子采用 Ascend **MIX AIC/AIV** 架构(`KERNEL_TYPE_MIX_AIC_1_2`,即 1 个 Cube 核搭配 2 个 Vector 核): + +- **Cube 侧(AIC)**:类 `MLAttentionDecoderAic`,负责两次矩阵乘 —— QK^T(mm1)与 PV(mm2)。 +- **Vector 侧(AIV)**:类 `MLADecoderAiv`,负责 Softmax 与 flash-attention 在线累加(online rescale)。**BF16 下不含反量化(DeQuant)/量化(Quant)步骤。** + +两侧通过 Workspace 上的中间 GM buffer 和跨核同步原语(FftsCrossCoreSync)协作,形成流水: +`QK^T(Cube) → Softmax(Vector) → PV(Cube) → Online Rescale 输出(Vector)`。 + +### 输入/输出 + +| 序号 | 名称 | 说明（BF16 场景） | +|------|------|------| +| 0 | query | Q 主体(**bf16**,hidden=576,含 nope 512 + rope 64) | +| 1 | queryRope | Q 的 rope 部分(bf16,hidden=64) | +| 2 | kvCache | KV Cache 主体(**bf16**,支持 ND / NZ 格式) | +| 3 | kvCacheRope | KV Cache 的 rope 部分(bf16) | +| 4 | block_tables | PagedAttention 块表 | +| 5 | contextLens | KV 序列长度 | +| 6 | mask | 注意力 mask | +| 7 | qSeqlen | Q 序列长度 | +| 8 | qkDescale | **BF16 场景不使用**（无 QK 反量化） | +| 9 | pvDescale | **BF16 场景不使用**（无 PV 反量化） | +| 10 | attenOut | 注意力输出(bf16) | +| 11 | lseOut | log-sum-exp 输出(ring 场景) | + +> BF16 场景下 `qkDescale`/`pvDescale` 两个量化 scale 输入不参与计算(全程无量化)。 + +### 数学定义与含义 + +本算子在 **decode 阶段**为每个 query token 计算一次标准的缩放点积注意力(scaled dot-product attention),但 K/V 来自 MLA 压缩的低秩隐空间,并按 PagedAttention 组织。 + +**1) 基础注意力公式** + +对第 `h` 个 head、当前 query 向量 `q_h`(与其历史 KV 序列 `K_h, V_h`,长度 = 上下文长度 `L`): + +``` +Attn_h = softmax( (q_h · K_hᵀ) / √d + mask ) · V_h +``` + +其中 `d` 为 head 维度,缩放系数 `tor = 1/√d`(host 侧算好写入 tiling 的 `TILING_TOR`)。 + +**2) MLA 的 rope 拼接** + +MLA 把 Q/K 拆成**压缩主体**(nope,hidden=512)与 **rope 位置编码部分**(hidden=64),QK^T 分数是两部分之和: + +``` +score = q_nope · k_nopeᵀ + q_rope · k_ropeᵀ +``` + +**BF16 下 nope 与 rope 不再拆成两条不同精度的 MMA**:因为主体本身就是浮点,rope 也是浮点,二者可以拼成 **hidden=576** 的统一 bf16 矩阵乘一次算出,无需像 INT8 那样把主体走 int8、rope 单独走 float 再相加(见 §5.3、§6.4)。 + +**3) BF16 下的等价计算（无量化）** + +Q、K、P 全程以 bf16 存储,矩阵乘在 float(fp32)域累加,直接得到结果,无 scale 还原: + +``` +score = (Q_bf16 · K_bf16ᵀ) # bf16 × bf16 → float(hidden=576 一次算完，含 rope) +P = softmax(score × tor + mask) # 概率 ∈ [0,1]，float 域 +Attn_h = (P_bf16 · V_bf16) # bf16 × bf16 → float，直接累加 +``` + +即 **没有 DeQuant(QK)/Requant(P)/DeQuant(PV) 三个量化点**;所有矩阵乘输入为 bf16、累加为 float，Softmax 全程 float,概率转回 bf16 仅用一次 `Cast`(不带 scale)。 + +**4) Flash-Attention 在线累加(online softmax)** + +与 INT8 完全一致。KV 按 block(block_size=64)逐段计算,采用 flash-attention 的在线归约。设历史最大值 `gm`、历史分母 `gl`、历史加权输出 `go`,新 block 的局部最大 `hm`、局部行和 `ll`、局部输出 `lo`: + +``` +m_new = max(gm, hm) +dm = exp(gm - m_new) # 历史项 rescale 因子 +gl = dm · gl + ll # 分母(归一化因子)累加 +go = dm · go + lo # 分子(∑ P·V)累加 +gm = m_new +``` + +全部 block 处理完后归一化输出: + +``` +attenOut_h = go / gl +lseOut_h = gm + log(gl) # ring/分布式场景需要的 log-sum-exp +``` + +--- + +## 2. 算子注册与数据类型 + +算子注册见 `op_host/multi_latent_attention_def.cpp`。**BF16 场景**的关键特征: + +- `query`、`kvCache` 数据类型为 `DT_BF16`; +- `kvCache` 的 Format 可为 `FORMAT_ND`(TILING_KEY 1)或 `FORMAT_FRACTAL_NZ`(TILING_KEY 17); +- `queryRope`/`kvCacheRope` 同为 bf16。 + +对应的模板实例化(见 `op_kernel/multi_latent_attention.cpp`): + +```cpp +// TILING_KEY 1: bf16(IN) + bf16(OUT), ND 格式 +MLAttentionDecoderAic +// TILING_KEY 17: bf16(IN) + bf16(OUT), NZ 格式 +MLAttentionDecoderAic +``` + +模板参数含义:输入类型 `__bf16`、rope 类型 `__bf16`、输出类型 `__bf16`、KV 类型 `__bf16`、输入格式 `ND_FORMAT/NZ_FORMAT`。**注意 INT8 用 5 个不同的类型参数(int8/half/half/int8),而 BF16 五个数据类型参数全部是 `__bf16`。** + +### 2.1 AttentionType 类型萃取(BF16 vs INT8） + +`AttentionType<>` 特化决定 mm1/mm2 的中间累加类型(见 `multi_latent_attention.h`): + +| 成员 | BF16 (= HALF) | INT8 | +|------|--------------|------| +| mm1OutputType / mm1CopyType | `float` | `int32_t` | +| mm2OutputType / mm2CopyType | `float` | `int32_t` | +| mmBiasType / mmScaleType | `float` | `float` | + +**BF16 与 fp16(HALF)的类型萃取完全相同**:两次矩阵乘的输出/搬运类型均为 `float`,即 bf16×bf16 累加到 float,不存在 int32 量化域。这是 BF16 与 INT8 在 kernel 层最根本的区别。 + +--- + +## 3. TilingKey 生成规则 + +见 `MLATiling()` → `GenTilingKey()`(`op_host/multi_latent_attention_tiling_impl.cpp`): + +```cpp +uint32_t dataType = static_cast(mmInfo.type); +uint32_t tilingKey = dataType + + (mmInfo.kNz << 4) // KV 是否 NZ 格式 + + (mmInfo.mtpTp1Flag << 2) // 是否 MTP/TP1 分支(numHeads==128) + + (param.isRing << 5); // 是否 ring attention +``` + +其中 `dataType` 取值(`GetTilingKeyTypeBase()`): + +| type 值 | 枚举 | 含义 | +|---------|------|------| +| 0 | TILING_HALF_DATA | fp16 | +| **1** | **TILING_BF16_DATA** | **bf16** | +| 2 | TILING_INT8_HALF_DATA | int8 输入 / fp16 输出 | +| 3 | TILING_INT8_BF16_DATA | int8 输入 / bf16 输出 | + +BF16 判定:当 `query` 为 bf16 时 `dataType = 1`。 + +**BF16 常见 TILING_KEY 组合**(`dataType=1`): + +| TILING_KEY | 组合 | 计算式 | +|------------|------|--------| +| 1 | bf16 + ND | `1` | +| 17 | bf16 + NZ | `1 + (1<<4)` | +| 5 | bf16 + ND + TP1 | `1 + (1<<2)` | +| 21 | bf16 + NZ + TP1 | `1 + (1<<4) + (1<<2)` | +| 33 | bf16 + ND + ring | `1 + (1<<5)` | +| 49 | bf16 + NZ + ring | `1 + (1<<4) + (1<<5)` | +| 37 / 53 | bf16 + ring + TP1(ND/NZ) | `+(1<<2)` | + +**与 INT8 的关键区别**:INT8 恒走 18/19(强制 NZ、不支持 TP1);**BF16 支持 ND 与 NZ 两种格式,且支持 MTP/TP1 分支**(`mtpTp1Flag = (numHeads == 128) && (type < 2)`,BF16 的 type=1 < 2 满足条件)。 + +--- + +## 4. Host 侧 Tiling 实现 + +Host Tiling 逻辑(`op_host/`)在 BF16 与 INT8 之间**基本一致**,仅在 workspace 各段的数据类型/字节大小与 hidden 维度上有差异。核心常量: + +| 常量 | 值 | 含义 | +|------|-----|------| +| TILING_HEAD_SIZE | 15 | tiling 头部字段数 | +| TILING_PARA_SIZE | 8 | 每个 batch 的字段数 | +| BATCH_MLA | 32 | 典型 batch | +| BLOCK_DIM_MLA | 20 | batch==32 时固定 20 个 Cube 核 | +| M_LIMIT | 128 | 单次处理的 M 上限 | + +### 4.1 核数与任务数 + +- `totalTaskNum = Σ qSeqLen`(decode 阶段每个 batch 的 qSeqLen 通常为 1,故 ≈ batch)。 +- `blockDim = GetCoreNumAic()`;当 `batch == 32` 时固定使用 20 个核。 +- 任务按 **round-robin** 方式在核间轮转分配。 + +### 4.2 tiling 数据布局 + +- **头部 15 项**:全局参数(numHead、hidden、tor、block_size、page 相关等)。 +- **每 batch 8 项**:`qSeqLen`、`kvSeqlen`,以及 query / kvCache / block_table 三组地址的高低 32 位。地址随 batch 逐个累加算出。 + +### 4.3 Workspace 分段 + +Workspace 划分为 6 段中间 GM buffer,BF16 场景各段按 **浮点(float/bf16)** 大小分配(INT8 场景 s_gm 走 int32、p_gm 走 int8): + +| 段 | 名称 | 用途 | BF16 类型 | +|----|------|------|-----------| +| 1 | s_gm | QK^T 分数 | float | +| 2 | s_rope_out_gm | (INT8 专用 rope 分数) | **BF16 不使用** | +| 3 | p_gm | Softmax 概率 | bf16(OUT_DTYPE) | +| 4 | o_tmp_gm | PV 中间输出 | float | +| 5 | go_gm | 在线累加输出 | float | +| 6 | tmp_gm | 临时 buffer | float | + +> BF16 下 `s_rope_out_gm` 段不参与计算 —— rope 已并入 hidden=576 的统一 MMA,分数直接落 `s_gm`。 + +--- + +## 5. Kernel 入口与 BF16 数据流 + +### 5.1 入口分发 + +`op_kernel/multi_latent_attention.cpp` 中 `extern "C"` 入口按 `TILING_KEY_IS` 分发。BF16 分支: + +```cpp +if (TILING_KEY_IS(1) || TILING_KEY_IS(17) /* ND / NZ */) { + // AIC: MLAttentionDecoderAic + // AIV: MLADecoderAiv +} +``` + +内核首先解析 6 段 workspace GM 地址,再按 Cube / Vector 角色进入各自主循环。 + +### 5.2 BF16 数据流(对比 INT8） + +``` + ┌──────────────── AIC (Cube) ────────────────┐ + query/kvCache ─▶│ QK^T: bf16 × bf16 → float (hidden=576, │─▶ s_gm(float) + (bf16) │ nope+rope 一次 MMA 算完) │ + └────────────────────────────────────────────┘ + │ QK_READY + ▼ + ┌──────────────── AIV (Vector) ──────────────┐ + s_gm(float) ───▶│ SoftmaxStage1: 读 float s_gm → ×tor → +mask │─▶ p_gm(bf16) + │ → rowmax → flash(max/dm) → exp │ + │ → Cast(float→bf16) ★无 Requant │ + └────────────────────────────────────────────┘ + │ SOFTMAX_READY + ▼ + ┌──────────────── AIC (Cube) ────────────────┐ + p_gm(bf16) ────▶│ PV: bf16 × bf16 → float │─▶ o_tmp_gm(float) + └────────────────────────────────────────────┘ + │ UPDATE_READY + ▼ + ┌──────────────── AIV (Vector) ──────────────┐ + o_tmp_gm ──────▶│ SoftmaxStage2: 读 float o_tmp → online │─▶ attenOut(bf16) + │ rescale(gl=dm·gl+ll / go=dm·go+lo) │─▶ lseOut(ring) + │ → 末 block go/gl → Cast ★无 DeQuant │ + └────────────────────────────────────────────┘ +``` + +**与 INT8 流程图的三处删减**: +1. QK^T 后**没有** DeQuant(int32×qkDescale→float);BF16 直接产出 float。 +2. SoftmaxStage1 后**没有** Requant(×1/127 → int8);BF16 直接 Cast float→bf16。 +3. PV 后**没有** DeQuant(int32×pvDescale→float);BF16 直接产出 float。 + +### 5.3 hidden 维度 + +BF16 场景 `hidden_size = 576`(nope 512 + rope 64,一并做 MMA);INT8 场景 `hidden_size = 512`(rope 64 单独走 float MMA 落 `s_rope_gm`)。`n_loop = (cur_kv_seqlen + pp_n - 1) / pp_n`。 + +--- + +## 6. AIC(Cube 侧)`MLAttentionDecoderAic` + +### 6.1 SetArgs / Run + +`Run()` 按 round-robin 领取本核负责的 (batch, task) 任务,循环调用 `InnerRunCubeMLA()`。 + +### 6.2 Q 地址与 L1 搬运 + +- 用 tiling 中每 batch 的地址高/低 32 位拼接出 64 位 Q / kvCache / block_table GM 地址。 +- 将 Q 搬入 L1。 + +### 6.3 n_loop 与 block_table 定位 + +按 KV 序列长度切成 `n_loop` 个 block,通过 `block_table` 定位每个 KV block 在 Cache 中的物理页。 + +### 6.4 CUBE1(QK^T,BF16 关键差异） + +BF16 走 **embed_split 5 段**(4×128 + 64 = 576),但**统一用 bf16 mmad 一次算完**: + +``` +非 INT8 / BF16 分支(multi_latent_attention.h L947-1003): + for idx in [0..4]: # 5 段 embed_split,累加到同一 L0C + mmad(bf16 × bf16 → float accumulate) + idx == 4 结束后: 一次 l0c_to_gm 写 s_gm(float) + ★无 rope 独立分支、★无 dequant +``` + +对比 INT8:INT8 在 `idx == 3` 时把主体分数(int32)写 `s_gm`,`idx == 4` 单独把 rope 部分以 float MMA 写 `s_rope_gm`,后续由 Vector 侧 DeQuant 再相加。**BF16 因全程浮点,rope 直接并入统一 MMA,少一次 GM 往返与一个量化点。** + +### 6.5 CUBE2(PV) + +用 `LoadDataWithTranspose` 把概率 `p_gm`(bf16)与 V(bf16)做矩阵乘,累加到 float 落 `o_tmp_gm`。 + +### 6.6 跨核同步 + +`FftsCrossCoreSync` 依次发出 `QK_READY → SOFTMAX_READY → UPDATE_READY`,与 Vector 侧握手。 + +--- + +## 7. AIV(Vector 侧)`MLADecoderAiv` + +`InnerRunVectorChange()` 将 head 按 `sub_block_idx`(0/1)分给两个 Vector 核各半,用 `n_loop + 1` 的软流水、`n_idx % 2` ping-pong buffer,使 Stage1 与 Stage2 错位一拍并行。 + +### 7.1 SoftmaxStage1(BF16 分支,`multi_latent_attention.h` L2517-2796） + +``` +BF16 (else 分支): + gm_to_ub : 直接把 float 的 s_gm 搬入 ls32_ubuf ★无 DeQuantPerHeadImpl、★无 s_rope_gm 相加 + mask : DataCopy + Cast(载入 mask) + muls(tor) : ls × tor(分 FLOAT_VECTOR_SIZE 段 + 尾段) + mask Add : + mask + ReduceMaxRepeatM : 行最大 lm + flash : n_idx!=0 → hm=max(lm,gm), dm=exp(gm-hm); else hm=lm; gm=hm + TensorSubValueRepeatM : ls - hm + exp_v : exp(ls - hm) + conv_v : float → OUT_DTYPE(bf16) ★无 QuantPerTokenImpl(不乘 1/127、不转 int8) + ub_to_gm : 写 p_gm(bf16) + ReduceSumRepeatM : 行和 ll +``` + +对比 INT8:INT8 分支先做 `DeQuantPerHeadImpl`(s_gm×qkDescale→float)、把 `s_rope_gm` 以 float 载入相加,末尾用 `QuantPerTokenImpl`(×1/127 转 int8)。**BF16 三处量化相关操作全部省去。** + +### 7.2 SoftmaxStage2MLAHeadLoop(BF16 分支,L2798-3157） + +``` +n_idx != 0: + gm_to_ub : 读 o_tmp_gm(float)作为 lo ★BF16 无 DeQuantPerHeadImpl + head_loop_idx==0: exp(dm); gl = dm·gl + ll + brcb dm → tv; go = go·dm_block(分段 mul_v); go = go + lo +n_idx == 0: + gl = ll; gm_to_ub 读 o_tmp_gm 作为 go ★BF16 无 DeQuant + +末 block (n_idx == n_loop-1): + gl_block brcb; go = go / gl_block(div_v 分段) + conv_v : go(float) → OUT_DTYPE(bf16) + DataCopyPad 写 o_gm(head_res / numhead_per_process / tail 三段) + IS_RING: ln(gl) + gm → lse → conv_v → ub_to_gm_align 写 lse_gm +否则 head_loop>1: ub_to_gm 写 go_gm +``` + +**BF16 全程 float 在线 rescale,无 PV 反量化步骤。** `process_row_num = 16` 分块,ring 场景额外输出 `lseOut`。 + +### 7.3 TP1 路径 + +`numHeads == 128` 时走 TP1 特化:`SoftmaxStage2MLAHeadLoopTP1` / `TailSoftmaxStage2MLAHeadLoopTP1` / `SoftmaxGatherTP1`,以及 `OnlineSoftmaxStage1`(`multi_latent_attention_npu.h`)。**该路径 BF16 可用,INT8 不支持。** + +--- + +## 8. 小结:BF16 与 INT8 的核心差异 + +| 维度 | BF16 | INT8 | +|------|------|------| +| mm1/mm2 累加类型 | float(bf16×bf16→float) | int32(int8×int8→int32) | +| 三量化点 | **全部无** | DeQuant(QK) / Requant(P,×1/127) / DeQuant(PV) | +| hidden_size | **576**(nope+rope 一体 MMA) | 512(rope 64 独立 float MMA) | +| s_rope_gm 段 | 不使用 | 使用(rope 分数) | +| QK^T | 5 段 embed_split 统一 mmad,一次写 s_gm | idx==3 写 s_gm、idx==4 rope 单独写 s_rope_gm | +| SoftmaxStage1 输出 | conv_v float→bf16 写 p_gm | Requant ×1/127 → int8 | +| SoftmaxStage2 | float online rescale,无 DeQuant | DeQuant PV×pvDescale + online rescale | +| 支持格式 | ND(1) / NZ(17) | 强制 NZ(18/19) | +| TP1 分支 | 支持(5/21/37/53) | 不支持 | +| qkDescale/pvDescale | 不使用 | 必需 | + +**一句话总结**:BF16 路径相较 INT8 —— **去掉全部三个量化点、hidden 统一为 576(rope 并入主 MMA)、矩阵乘直接 bf16×bf16→float 累加、Softmax 全程 float,概率仅用一次不带 scale 的 Cast 转回 bf16**。数学上与标准缩放点积注意力 + flash online softmax 完全等价,计算链路比 INT8 更短更直接。 \ No newline at end of file diff --git a/xllm_ops/multi_latent_attention/multi_latent_attention_desc.md b/xllm_ops/multi_latent_attention/multi_latent_attention_desc.md new file mode 100644 index 0000000..f5fa712 --- /dev/null +++ b/xllm_ops/multi_latent_attention/multi_latent_attention_desc.md @@ -0,0 +1,524 @@ +# Multi-Latent Attention (MLA) 算子实现分析 + +> 本文档分析 `xllm_ops/multi_latent_attention` 算子在 Ascend AscendC 平台上的实现。 +> - **Host 侧**:聚焦 Tiling 切分策略(tiling 参数、核数计算、任务数、多核分配)。 +> - **Kernel 侧**:聚焦 **INT8 数据类型**的实现(业务处理流程、函数逻辑、数据获取、地址计算)。 + +--- + +## 1. 算子概述 + +Multi-Latent Attention(MLA)是 DeepSeek 系列模型使用的注意力机制,核心特点是把 KV Cache 压缩到一个低秩的隐空间(latent),从而大幅降低 KV Cache 显存占用。本算子实现的是 **decode(增量推理)阶段**的 MLA,基于 **PagedAttention** 的 block_table 机制管理 KV Cache。 + +算子采用 Ascend **MIX AIC/AIV** 架构(`KERNEL_TYPE_MIX_AIC_1_2`,即 1 个 Cube 核搭配 2 个 Vector 核): + +- **Cube 侧(AIC)**:类 `MLAttentionDecoderAic`,负责两次矩阵乘 —— QK^T(mm1)与 PV(mm2)。 +- **Vector 侧(AIV)**:类 `MLADecoderAiv`,负责反量化(DeQuant)、Softmax、量化(Quant)与 flash-attention 在线累加(online rescale)。 + +两侧通过 Workspace 上的中间 GM buffer 和跨核同步原语(FftsCrossCoreSync)协作,形成流水: +`QK^T(Cube) → Softmax(Vector) → PV(Cube) → Online Rescale 输出(Vector)`。 + +### 输入/输出 + +| 序号 | 名称 | 说明 | +|------|------|------| +| 0 | query | Q 主体(INT8 场景为 int8,hidden=512) | +| 1 | queryRope | Q 的 rope 部分(float/half/bf16,hidden=64) | +| 2 | kvCache | KV Cache 主体(INT8 场景 int8,NZ 格式) | +| 3 | kvCacheRope | KV Cache 的 rope 部分 | +| 4 | block_tables | PagedAttention 块表 | +| 5 | contextLens | KV 序列长度 | +| 6 | mask | 注意力 mask | +| 7 | qSeqlen | Q 序列长度 | +| 8 | qkDescale | QK^T 反量化 scale(per-head float) | +| 9 | pvDescale | PV 反量化 scale(per-head float) | +| 10 | attenOut | 注意力输出 | +| 11 | lseOut | log-sum-exp 输出(ring 场景) | + +### 数学定义与含义 + +本算子在 **decode 阶段**为每个 query token 计算一次标准的缩放点积注意力(scaled dot-product attention),但 K/V 来自 MLA 压缩的低秩隐空间,并按 PagedAttention 组织。 + +**1) 基础注意力公式** + +对第 `h` 个 head、当前 query 向量 `q_h`(与其历史 KV 序列 `K_h, V_h`,长度 = 上下文长度 `L`): + +``` +Attn_h = softmax( (q_h · K_hᵀ) / √d + mask ) · V_h +``` + +其中 `d` 为 head 维度,缩放系数 `tor = 1/√d`(host 侧算好写入 tiling 的 `TILING_TOR`)。 + +**2) MLA 的 rope 拼接** + +MLA 把 Q/K 拆成**压缩主体**(nope,hidden=512)与 **rope 位置编码部分**(hidden=64),QK^T 分数是两部分之和: + +``` +score = q_nope · k_nopeᵀ + q_rope · k_ropeᵀ +``` + +因此 kernel 里主体走 int8 量化 MMA,rope 部分单独走 float MMA,二者在 softmax 前相加(见 §5.3、§7.3)。 + +**3) INT8 量化下的等价计算** + +主体 Q、K、P 均以 int8 存储,矩阵乘在 int32 域累加,再用 per-head/per-token scale 还原: + +``` +q_nope · k_nopeᵀ ≈ (Q_int8 · K_int8ᵀ) × qkDescale # int32 → float,per-head 反量化 +score = 上式 + q_rope · k_ropeᵀ # rope 恒为 float +P = softmax(score × tor + mask) # 概率 ∈ [0,1] +P_int8 = round(P × 127) # per-token 量化(scale = 1/127) +Attn_h ≈ (P_int8 · V_int8) × pvDescale × (1/127) # int32 → float,per-head 反量化 +``` + +即三个量化点:**DeQuant(QK)→ Requant(P)→ DeQuant(PV)**;rope 分支始终保持浮点精度。 + +**4) Flash-Attention 在线累加(online softmax)** + +由于 KV 按 block(block_size=64)逐段计算,采用 flash-attention 的在线归约,避免一次性物化整条注意力矩阵。设历史最大值 `gm`、历史分母 `gl`、历史加权输出 `go`,新 block 的局部最大 `hm`、局部行和 `ll`、局部输出 `lo`: + +``` +m_new = max(gm, hm) +dm = exp(gm - m_new) # 历史项 rescale 因子 +gl = dm · gl + ll # 分母(归一化因子)累加 +go = dm · go + lo # 分子(∑ P·V)累加 +gm = m_new +``` + +全部 block 处理完后归一化输出: + +``` +attenOut_h = go / gl +lseOut_h = gm + log(gl) # ring/分布式场景需要的 log-sum-exp +``` + +--- + +## 2. 算子注册与数据类型 + +算子注册见 `op_host/multi_latent_attention_def.cpp`。其中 query/kvCache 支持多种数据类型组合,**INT8 场景**的关键特征: + +- `query`、`kvCache` 数据类型为 `DT_INT8`; +- `kvCache` 的 Format 为 `FORMAT_FRACTAL_NZ`(NZ 格式); +- `queryRope`/`kvCacheRope` 仍为浮点(fp16 或 bf16),rope 部分不量化。 + +对应的模板实例化(见 `op_kernel/multi_latent_attention.cpp`): + +```cpp +// TILING_KEY 18: int8(IN) + fp16(OUT) +MLAttentionDecoderAic +// TILING_KEY 19: int8(IN) + bf16(OUT) +MLAttentionDecoderAic +``` + +模板参数含义:输入类型 `int8_t`、输出类型 `half/__bf16`、中间/bias 类型、量化类型 `int8_t`、输入格式 `NZ_FORMAT`。 + +--- + +## 3. TilingKey 生成规则 + +见 `MLATiling()` → `GenTilingKey()`(`op_host/multi_latent_attention_tiling_impl.cpp`): + +```cpp +uint32_t dataType = static_cast(mmInfo.type); +uint32_t tilingKey = dataType + + (mmInfo.kNz << 4) // KV 是否 NZ 格式 + + (mmInfo.mtpTp1Flag << 2) // 是否 MTP/TP1 分支(numHeads==128) + + (param.isRing << 5); // 是否 ring attention +``` + +其中 `dataType` 取值(`GetTilingKeyTypeBase()`): + +| type 值 | 枚举 | 含义 | +|---------|------|------| +| 0 | TILING_HALF_DATA | fp16 | +| 1 | TILING_BF16_DATA | bf16 | +| 2 | TILING_INT8_HALF_DATA | int8 输入 / fp16 输出 | +| 3 | TILING_INT8_BF16_DATA | int8 输入 / bf16 输出 | + +INT8 判定:当 `query` 不是 bf16/fp16 时进入 INT8 分支;再看 `queryRope` 是 fp16(→ type=2)还是 bf16(→ type=3)。 + +由于 INT8 的 KV Cache 强制 NZ(`kNz=1`,`<<4` 即 +16),最终 **INT8 走 TILING_KEY 18(fp16 输出)/ 19(bf16 输出)**: +- `2 + (1<<4) = 18` +- `3 + (1<<4) = 19` + +INT8 不支持 MTP/TP1(`mtpTp1Flag` 要求 `type < 2`),因此 INT8 恒走非 TP1 的 `Run()` 路径。 + +--- + +## 4. Host 侧 Tiling 切分策略 + +Tiling 的入口是 `MLATiling()`,主要逻辑分布在: +- `op_host/multi_latent_attention_tiling_impl.cpp` —— 主入口、信息采集、TilingKey、Workspace 计算。 +- `op_host/multi_latent_attention_tiling_dependency.cpp` —— tiling 参数填充、核数与任务分配。 + +### 4.1 关键常量 + +| 常量 | 值 | 含义 | +|------|-----|------| +| `TILING_HEAD_SIZE` | 15 | tiling 头部(公共参数)占用的 uint32 个数 | +| `TILING_PARA_SIZE` | 8 | 每个 batch 任务参数占用的 uint32 个数 | +| `TILING_PARA_SIZE_TP1` | 4 | TP1 分支每个 task 的参数个数 | +| `BATCH_MLA` | 32 | 特殊 batch 数(触发固定核数) | +| `BLOCK_DIM_MLA` | 20 | batch==32 时固定使用的核数 | +| `M_LIMIT` | 128 | numHeads==128 时走 MTP/TP1 分支 | +| `PP_MM` | {16,32,...,128} | M 方向分块候选 | +| `QN_TILE_LIST` | {128,64,32,16,8,1} | Q head 方向的分块候选 | + +### 4.2 信息采集(GetMLANdInfo) + +从 `TilingContext` 提取形状与属性: +- **NZ 判定**:`kNz = (kvCache 末维 == 16 或 32) ? 1 : 0`。INT8 KV Cache 为 NZ,`kNz=1`。 +- **embeddingSize / blockSize**:NZ 格式下 `embeddingSize = dim3 * dim1`,`blockSize = dim2`;ND 格式下取原始维度。 +- **batch** = `kvSeqLen.size()`(KV 序列条数)。 +- **numHeads** = 属性 `headSize`;**kvHeads** = `kvHead`(≤0 则等于 numHeads)。 +- **mtpTp1Flag** = `(numHeads == 128) && (type < 2)` —— INT8 恒为 false。 + +### 4.3 任务数(totalTaskNum)计算 + +```cpp +if (mmInfo.qSeqLen != nullptr) { + // 所有 batch 的 qSeqLen 之和 + mmInfo.totalTaskNum = accumulate(qSeqLen, qSeqLen + batch, 0); +} else { + mmInfo.totalTaskNum = batch; // decode 每 batch 一个 task +} +``` + +decode 场景每个 batch 的 qSeqLen 通常为 1,因此 **totalTaskNum 一般等于 batch**。该值写入 tiling 头部 `TILING_TASK_NUM` 供 kernel 侧划分 process。 + +### 4.4 核数(blockDim)计算 + +核数计算见 `MLATiling()` 与 `GetMLATilingParam()`: + +```cpp +auto blockDim = ascendcPlatform.GetCoreNumAic(); // 默认取平台 AIC 核数 +... +// 非 TP1 分支 +blockDim = mmInfo.batch == BATCH_MLA ? BLOCK_DIM_MLA : blockDim; +``` + +- 默认 `blockDim` 取硬件 **AIC 核数**(`GetCoreNumAic()`)。 +- **特殊优化**:当 `batch == 32` 时,固定使用 `BLOCK_DIM_MLA = 20` 个核 —— 针对该 batch 规模做过负载均衡调优。 +- 最终 `context->SetBlockDim(blockDim)` 下发。由于是 MIX 架构,该 blockDim 表示 AIC 核数,对应 2×blockDim 个 AIV 核。 + +### 4.5 Tiling 参数布局 + +Tiling data 的整体布局(`tilingParam` 指针): + +``` +[0..5] : 6 个 uint64 的 workspace 段大小(占 6*2 个 uint32) +[6*2..] : tiling 头部(TILING_HEAD_SIZE=15 个 uint32) + + 每 batch 参数(TILING_PARA_SIZE=8 个 uint32)× batch +``` + +**tiling 头部字段**(`GetTilingHead()`,下标见 `_dependency.cpp`): + +| 下标 | 字段 | 含义 | +|------|------|------| +| 0 | TILING_BATCH | batch 数 | +| 1 | TILING_NUMHEADS | numHeads | +| 2 | TILING_HEADDIM | embeddingSize | +| 3 | TILING_NUMBLOKS | numBlocks | +| 4 | TILING_BLOCKSIZE | blockSize | +| 5 | TILING_MAXBLOCKS | maxNumBlocksPerQuery | +| 6 | TILING_TOR | 缩放系数 tor(float 位模式) | +| 7 | TILING_KVHEADS | kvHeads | +| 8 | TILING_HEADSIZE | =15(头部大小) | +| 9 | TILING_PARASIZE | 每 task 参数大小(8 或 TP1 的 4) | +| 12 | TILING_MASK_TYPE_ND | maskType | +| 13 | TILING_TASK_NUM | totalTaskNum | +| 14 | TILING_MAX_KV_SEQ_LEN | maxKVseqlen | + +**每 batch 参数字段**(`GetNdMLATiling()` + `GetAddrOffsetMLA()`,偏移 `tilingOffset = 15 + 8*seqIdx`): + +| 偏移 | 字段 | 含义 | +|------|------|------| +| +0 | qSeqLen | 该 batch 的 Q 序列长度 | +| +1 | kvSeqlen | 该 batch 的 KV 序列长度 | +| +2/+3 | addrQSeqOffset 高/低 32 位 | Q/O 的累积地址偏移(64 位拆分) | +| +4/+5 | addrOSeqOffset 高/低 32 位 | 输出地址偏移 | +| +6/+7 | addrMaskOffset 高/低 32 位 | mask 地址偏移 | + +地址偏移**逐 batch 累加**: +```cpp +addrQSeqOffset += numHeads * qSeqLen; +addrOSeqOffset += numHeads * embeddingSize * qSeqLen; +addrMaskOffset += qSeqLen * maxKVseqlen; +``` +kernel 侧读取时把高低 32 位重新拼成 64 位地址,再乘以每 head 的 element 数得到实际 GM 偏移。 + +### 4.6 多核任务分配 + +- Host 侧仅确定 **核数(blockDim)** 与 **总任务数(totalTaskNum)** 以及每个 batch 的参数/地址偏移。 +- **实际的 task→core 映射在 kernel 侧动态完成**:每个核用自身 `block_idx`(0..blockDim-1)以 `blockDim` 为步长循环领取 process(见 §6 的 `Run()`),即典型的 **round-robin 静态均分**。 +- INT8 场景总 process = `q_block(每 batch 内 head 分块数) × batch`,由各 AIC/AIV 核以 `block_idx` 起步、步长 `blockDim` 遍历。 + +### 4.7 Workspace 切分 + +`MLATiling()` 计算 6 段 workspace(`workspaceParam[0..5]`),INT8(isQuant)与浮点分配不同: + +| 段 | 变量 | INT8(isQuant) | 浮点 | +|----|------|---------------|------| +| 0 | s_gm | basicWorkSpaceFloat | float×2 | +| 1 | s_rope_out_gm | basicWorkSpaceFloat | 512 | +| 2 | p_gm | basicWorkSpaceInt8 | half×2 | +| 3 | o_tmp_gm | basicWorkSpaceInt8×2 | float×2 | +| 4 | go_gm | basicWorkSpaceFloat | float | +| 5 | tmp_gm | tailWorkSpaceFloat | float | + +其中 `basicWorkSpace* = blockDim * WORKSPACE_BLOCK_SIZE_DB * dataLen`,即按核数 double-buffer 分配。INT8 的 p_gm/o_tmp_gm 用 int(int8/int32)存储,总 usrSize 再加系统 workspace。 + +--- + +## 5. Kernel 入口与 INT8 整体业务流程 + +### 5.1 Kernel 入口 + +入口函数 `multi_latent_attention()`(`op_kernel/multi_latent_attention.cpp`)按顺序完成: + +1. 声明 MIX 任务类型 `KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2)`。 +2. 从 workspace 头部解析 6 段中间 GM 地址(段大小由 host 的 `workspaceParam[0..5]` 给出): + +```cpp +GM_ADDR s_gm = usrWorkspace; // QK^T 结果(INT8 为 int32) +GM_ADDR s_rope_out_gm= s_gm + workspaceParam[0]; // QK rope 部分(float) +GM_ADDR p_gm = s_rope_out_gm + workspaceParam[1]; // softmax 概率(INT8 为 int8) +GM_ADDR o_tmp_gm = p_gm + workspaceParam[2]; // PV 结果(INT8 为 int32) +GM_ADDR go_gm = o_tmp_gm + workspaceParam[3]; // online 累加输出(float) +GM_ADDR tmp_gm = go_gm + workspaceParam[4]; // 临时 buffer +``` + +3. tiling 参数区从 `tiling + sizeof(uint64)*6` 开始(跳过 6 段 workspace 大小)。 +4. 按 `TILING_KEY` 分发。INT8 走 **18(int8→fp16)/ 19(int8→bf16)**: + - `__DAV_C220_CUBE__`(AIC):实例化 `MLAttentionDecoderAic<...>`,调 `SetArgs()` → `Run()`。 + - `__DAV_C220_VEC__`(AIV):实例化 `MLADecoderAiv<...>`,调 `SetArgs()` → `Run()`。 + +### 5.2 INT8 类型映射(AttentionType) + +INT8 场景下 mm1/mm2 的关键类型: +- **mm1/mm2 Output/CopyType = `int32_t`**:两次矩阵乘(int8×int8)累加结果均为 int32。 +- **mmBias/mmScaleType = `float`**:反量化 scale 为 float。 +- 输入 `IN_DTYPE = int8_t`,rope 部分 `IN_ROPE_DTYPE = half/bf16`(不量化)。 + +### 5.3 INT8 完整数据流 + +一次注意力计算(单个 KV block)的 INT8 数据流如下: + +``` + ┌─────────────── AIC (Cube) ───────────────┐ + Q(int8,512) ─┐ │ CUBE1: QK^T │ + K(int8,512) ─┼─ mmad ──►│ int8 × int8 → int32 ──► s_gm (int32) │ + Qrope(fp) ───┤ │ rope: float MMA ──► s_rope_gm (float)│ + Krope(fp) ───┘ └───────────────────────────────────────────┘ + │ FftsCrossCoreSync(QK_READY) + ▼ + ┌─────────────── AIV (Vector) ──────────────┐ + │ SoftmaxStage1: │ + │ DeQuantPerHead(s_gm × qk descale) → float │ + │ + s_rope_gm(float) → muls(tor) → mask │ + │ → rowmax → flash max/dm → exp │ + │ → QuantPerToken(× 1/127 scale) → int8 │ + │ → p_gm (int8) │ + └───────────────────────────────────────────┘ + │ FftsCrossCoreSync(SOFTMAX_READY) + ▼ + ┌─────────────── AIC (Cube) ───────────────┐ + p(int8) ─┐ │ CUBE2: PV │ + K^T(int8)┼─ mmad ──────►│ int8 × int8 → int32 ──► o_tmp_gm (int32) │ + └ │ (K 用 LoadDataWithTranspose 转置) │ + └───────────────────────────────────────────┘ + │ FftsCrossCoreSync(UPDATE_READY) + ▼ + ┌─────────────── AIV (Vector) ──────────────┐ + │ SoftmaxStage2MLAHeadLoop: │ + │ DeQuant(o_tmp_gm × pv descale) → float │ + │ online rescale: │ + │ dm = exp(gm - hm) │ + │ gl = dm*gl + ll ; go = go*dm + lo │ + │ 最后一个 block: go/gl → 输出 o_gm │ + └───────────────────────────────────────────┘ +``` + +**INT8 相比浮点的三个量化点**: +1. **DeQuant(QK)**:CUBE1 产出的 int32 乘 `qkDescale`(per-head float)还原为 float。 +2. **Requant(P)**:softmax 概率 P 用 per-token scale(`quantMax=1/127`)量化回 int8,供 CUBE2 用 int8 做 PV。 +3. **DeQuant(PV)**:CUBE2 产出的 int32 乘 `pvDescale`(per-head float)还原为 float。 + +**rope 部分始终走 float 独立 MMA**,不参与量化,在 SoftmaxStage1 中与 DeQuant 后的主体结果相加。 + +hidden_size:INT8 主体 =512(rope 的 64 单独处理),浮点场景为 576(512+64)。 + +--- + +## 6. AIC(Cube 侧)函数细节 —— MLAttentionDecoderAic + +Cube 侧类 `MLAttentionDecoderAic` 负责两次矩阵乘。核心执行流程:`SetArgs()` 保存 GM 指针与参数 → `Run()` 以 `block_idx` 为起点、`blockDim` 为步长 round-robin 领取 process → 每个 process 调 `InnerRunCubeMLA()` 完成 QK^T(CUBE1)与 PV(CUBE2)。 + +### 6.1 SetArgs / Run + +- **SetArgs**:保存 q_gm、q_rope_gm、ctkv_gm、ctkv_rope_gm、block_tables_gm、o_gm 以及 6 段 workspace GM 指针;从 `tiling_para_gm` 读头部公共参数(batch、numHeads、embeddingSize、blockSize、maxNumBlocksPerQuery、tor 等)。 +- **Run**:以 `block_idx`(核号)为起点、`blockDim` 为步长遍历 process。每个 process 用其 batch 的 `offset_tiling = TILING_HEAD_SIZE + TILING_PARA_SIZE * seqIdx` 定位到该 batch 参数,调用 `InnerRunCubeMLA()`。 + +### 6.2 InnerRunCubeMLA —— QK^T 与 PV + +单次处理一个 process(某 batch 的一段 head)。 + +#### (1) Q 地址计算 + +从 tiling 参数区读 Q 的 64 位地址偏移(高低 32 位拼接),再换算成 element 偏移: + +```cpp +uint64_t addr_q_scalar = ((uint64_t)addr_q_high32 << 32) | addr_q_low32; +uint64_t q_offset = addr_q_scalar * 512 + start_head * 512; // INT8 主体 hidden=512 +uint64_t q_rope_offset = addr_q_scalar * 64 + start_head * 64; // rope hidden=64 +``` + +INT8 主体 `hidden_size = 512`(浮点为 576),rope 部分独立按 64 计算偏移;`start_head` 为该 process 负责的起始 head。 + +#### (2) Q 搬入 L1 + +- `cur_q_seqlen == 1`(纯 decode):用 `gm_to_l1` 直接搬入 L1。 +- 否则用 `Nd2NzParams` 做 ND→NZ 转换搬入;head 数超过阈值时逐 seqlen 分批搬。 + +#### (3) n_loop 循环 —— 遍历 KV block + +对每个 KV block: +1. **block_table 定位**:通过 `block_tables_gm` 找到该逻辑 block 对应 KV Cache 的物理 block 号,算出 `kv_offset`。 +2. **K / K_rope 搬入 L1**:INT8 KV 为 NZ 格式,走 NZ→NZ 的 `gm_to_l1` 搬运。 + +#### (4) CUBE1:QK^T(embed_split 分段) + +hidden 128 方向切 5 段(前 4 段各 128,第 5 段为 rope 的 64),逐段: +- L1→L0A(Q)、L1→L0B(K); +- **INT8**:`mmad<..., int8_t, int8_t, int32_t, false>`(int8×int8→int32)累加到 `mm1_l0c`,`init` 标志在 `embed_split_idx == 0` 时置位。 + +```cpp +if constexpr (tilingKeyType == TILING_INT8_DATA) { + mmad<..., IN_DTYPE, IN_DTYPE, mm1OutputType, false>( // int8×int8→int32 + mm1_l0c, l0a, l0b, m, qk_round_n, embed_split_size, embed_split_idx == 0); +} +``` + +- **INT8 特殊分段**:`embed_split_idx == 3` 时把当前 int32 累加结果 `l0c_to_gm` 写到 `s_gm`;第 5 段(`idx == 4`)单独做 **rope 部分的 float MMA**,结果 `l0c_to_gm` 写到 `s_rope_gm`(float)。 + +```cpp +l0c_to_gm<..., mm1CopyType, mm1OutputType>(s_gm_tensor[...], mm1_l0c, ...); // int32 主体 +mmad<..., IN_ROPE_DTYPE, IN_ROPE_DTYPE, float, false>(...); // rope float +l0c_to_gm<..., float, float>(s_rope_gm_tensor[...], ...); +``` + +#### (5) CUBE2:PV(n_idx != 0 时) + +第一个 KV block 之后开始做上一 block 的 PV(与当前 block 的 QK 流水重叠): +- **K 转置**:用 `LoadDataWithTranspose` 把 K 从 L1 转置进 L0B; +- **P 搬入**:softmax 输出的 p 从 `p_gm`(int8)搬进 L0A(NZ→ZZ); +- **mmad**:`int8 × int8 → int32`,结果 `l0c_to_gm` 写到 `o_tmp_gm`(int32)。 + +```cpp +mmad<..., IN_DTYPE, IN_DTYPE, mm2OutputType, false>( // int8×int8→int32 + mm2_l0c, l0a_p, l0b_kT, m, embed_split_size, qk_n_2, 1); +l0c_to_gm<..., mm2CopyType, mm2OutputType>(o_tmp_gm_tensor[...], mm2_l0c, ...); +``` + +#### (6) 同步 + +Cube 与 Vector 通过 `FftsCrossCoreSync` 跨核同步:CUBE1 完成发 `QK_READY_DECODER`;等 Vector 的 `SOFTMAX_READY_DECODER` 后才做 CUBE2;PV 完成发 `UPDATE_READY_DECODER`。核内用 `SET_FLAG/WAIT_FLAG`(MTE2/MTE1/M/FIX)与 `PIPE_BARRIER` 保证 L1/L0A/L0B/L0C 的 ping-pong(16384 偏移)读写顺序。 + +--- + +## 7. AIV(Vector 侧)函数细节 —— MLADecoderAiv + +Vector 侧类 `MLADecoderAiv` 承担反量化、Softmax、量化与 flash-attention 在线累加,是 INT8 精度处理的核心。两个 Vector 核(`sub_block_idx` = 0/1)各处理一半 head。 + +### 7.1 InnerRunVectorChange —— AIV 主控 + +```cpp +uint32_t sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num/2) : cur_head_num/2; +uint32_t sub_m = sub_head_num * cur_q_seqlen; +o_offset = addr_o_scalar + start_head*embedding_size + sub_block_idx*cur_head_num/2*embedding_size; +``` + +- **head 切分**:`sub_block_idx` 0/1 各处理 `cur_head_num/2` 个 head;`sub_m = sub_head_num * cur_q_seqlen` 是本核处理的行数。 +- **o_offset**:输出地址按 `start_head` 与 `sub_block_idx` 偏移,两核写不同 head 区间。 +- **n_loop 循环**(按 `block_size = 64` 切 KV,循环 `n_loop + 1` 次做软件流水): + +```cpp +for (n_idx = 0; n_idx < n_loop + 1; n_idx++) { + if (n_idx != n_loop) { // Stage1:当前 block 的 softmax + WaitFlagDev(QK_READY_DECODER); // 等 Cube 的 QK^T 完成 + WAIT_FLAG(MTE3, MTE2, EVENT_ID3); + SoftmaxStage1(p_gm[...], s_gm[...], s_rope_gm[...], mask_gm[...], ...); // ping-pong(n_idx%2) + FftsCrossCoreSync(SOFTMAX_READY_DECODER); // 通知 Cube 做 PV + SET_FLAG(MTE3, MTE2, EVENT_ID3); + } + if (n_idx != 0) { // Stage2:上一 block 的 online rescale + WaitFlagDev(UPDATE_READY_DECODER); // 等 Cube 的 PV 完成 + uint32_t head_loop = (sub_m + process_row_num - 1) / process_row_num; // process_row_num=16 + for (uint32_t hl = 0; hl < head_loop; ++hl) { + SoftmaxStage2MLAHeadLoop(o_tmp_gm[...], go_gm[...], o_gm[o_offset + ...], ...); + } + } +} +``` + +- **ping-pong**:`n_idx % 2` 交替使用不同的 ubuf/gm 偏移(`dm32_ubuf`/`ll_ubuf`/`pm32_ubuf` 两组),使相邻 block 的 Stage1/Stage2 可重叠。 +- Stage1 与 Stage2 在同一次循环里错位一拍:`n_idx` 做当前 block 的 Stage1,同时做上一 block(`n_idx-1`)的 Stage2。 + +### 7.2 DeQuantPerHeadImpl —— QK 反量化 + +把 CUBE1 的 int32 结果按 per-head 的 `qkDescale` 还原为 float: + +```cpp +// 1. descale 搬入 ub;int32 结果搬入 ub +// 2. Cast int32 → float +Cast(float_ub, int32_ub, RoundMode::CAST_NONE, ...); +// 3. 逐 head 乘 descale(broadcast 到该 head 的所有列) +TensorMulRepeatM(float_ub, float_ub, descale_ub, ...); +``` + +每个 head 有独立 descale,因此按 head 循环做 broadcast 乘法。 + +### 7.3 SoftmaxStage1 —— DeQuant + rope + flash softmax + Requant + +单个 KV block 的 softmax,输出量化后的 int8 概率 P: + +1. **DeQuant + rope 合并**:`DeQuantPerHeadImpl(s_gm × qkDescale)` 得主体 float,再加上 `s_rope_gm`(float,rope 部分)。 +2. **缩放 + mask**:`Muls(x, tor)`(tor 为 1/√d 缩放系数),再叠加 `mask_gm`。 +3. **行最大 + flash 更新**:`ReduceMax` 求当前 block 行最大 `hm`;与历史最大 `gm` 比较更新,`dm = exp(gm - hm)` 作为历史部分的 rescale 因子。 +4. **exp**:`Exp(p, x - hm)` 得未归一化概率;`ll = rowsum(p)` 为当前 block 行和。 +5. **Requant(P)**:调 `QuantPerTokenImpl` 把 float 概率按 per-token scale(`quantMax = 1/127`)量化成 int8,写入 `p_gm` 供 CUBE2 使用。 + +### 7.4 QuantPerTokenImpl —— P 的 per-token 量化 + +```cpp +// scale = 1/127(per-token);float → int8 +Muls(x, x, scale); // 乘 1/127 +Cast(half_ub, x, RoundMode::CAST_NONE, ...); // float → half +Cast(int8_ub, half_ub, RoundMode::CAST_RINT, ...); // half → int8(四舍五入) +``` + +概率恒为正且 ≤1,用固定 `1/127` scale 映射到 int8 范围,再由 CUBE2 用 int8×int8 做 PV。 + +### 7.5 SoftmaxStage2MLAHeadLoop —— PV 反量化 + online rescale + +对 CUBE2 的 int32 PV 结果做反量化并做 flash-attention 在线累加: + +1. **DeQuant(PV)**:`o_tmp_gm`(int32)Cast→float 后乘 per-head `pvDescale`。 +2. **online rescale**(flash-attention 累加): + +``` +dm = exp(gm - hm) // 历史 rescale 因子(Stage1 已算) +gl = dm * gl + ll // 更新分母(行和) +go = go * dm + lo // 更新分子(加权 V 累加) +``` + +3. **收尾输出**:遍历到最后一个 block 后,`o = go / gl` 得到归一化注意力输出,`Cast` 成输出类型(fp16/bf16)写到 `o_gm[o_offset]`;`head_loop` 按 `process_row_num = 16` 行分块处理,避免 UB 溢出。ring 场景另写 `lseOut`。 + +--- + +## 8. 小结 + +- **架构**:MIX 1 Cube + 2 Vector,Cube 管两次 matmul,Vector 管量化/softmax/累加,靠 workspace GM + FftsCrossCoreSync 流水协作。 +- **Host Tiling**:核数默认 `GetCoreNumAic()`,`batch==32` 固定 20 核;任务数 = ΣqSeqLen(decode≈batch);tiling 由 15 项头部 + 每 batch 8 项(含地址高低 32 位)构成;task→core 在 kernel 侧 round-robin。 +- **INT8 全链路**:QK(int8×int8→int32)→ DeQuant(×qkDescale)+ rope(float)→ flash softmax → Requant(×1/127→int8)→ PV(int8×int8→int32)→ DeQuant(×pvDescale)→ online rescale → 输出。三个量化点 + rope 独立 float 路径是 INT8 与浮点实现的核心差异。 \ No newline at end of file From 197b1fcb87c4c7a83c65831a9baec03ba95ced1b Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Thu, 6 Aug 2026 17:40:00 +0800 Subject: [PATCH 07/24] feat: you hua mla --- .../op_kernel/multi_latent_attention.h | 923 ++++++++++-------- 1 file changed, 514 insertions(+), 409 deletions(-) diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h index 5bcb2ca..7f9bb38 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h @@ -690,33 +690,91 @@ class MLAttentionDecoderAic { } private: - __aicore__ __attribute__((always_inline)) inline void InnerRunCubeMLA(uint32_t cur_batch, uint32_t start_head, uint32_t cur_head_num, + // ====== MLA Cube refactor: context struct + sub-functions ====== + + struct MLAContext { + // input params + uint32_t cur_batch; + uint32_t start_head; + uint32_t cur_head_num; + uint32_t start_kv; + uint32_t cur_q_seqlen; + uint32_t cur_kv_seqlen; + uint32_t offset_tiling; + + // Q address + uint64_t q_offset; + uint64_t q_rope_offset; + + // loop & size + uint32_t pp_n_scalar; + uint32_t sub_n_loop; + uint32_t n_loop; + + // QK dims + uint32_t qk_n; + uint32_t qk_round_n; + uint32_t qk_n_2; + uint32_t qk_round_n_2; + uint32_t qk_round_n_l1; + uint32_t qk_round_n_2_l1; + + // hidden size + uint64_t hidden_size; + + // K round + uint64_t k_round_n; + + // row info + uint32_t row_num; + // m is class member + }; + + __aicore__ __attribute__((always_inline)) inline void InitMLAContext( + MLAContext &ctx, uint32_t cur_batch, uint32_t start_head, uint32_t cur_head_num, uint32_t start_kv, uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t offset_tiling) { + ctx.cur_batch = cur_batch; + ctx.start_head = start_head; + ctx.cur_head_num = cur_head_num; + ctx.start_kv = start_kv; + ctx.cur_q_seqlen = cur_q_seqlen; + ctx.cur_kv_seqlen = cur_kv_seqlen; + ctx.offset_tiling = offset_tiling; + uint32_t addr_q_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); uint32_t addr_q_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 3 + offset_tiling)); uint64_t addr_q_scalar = (uint64_t)(((uint64_t)addr_q_high32) << 32 | addr_q_loww32); - uint64_t q_offset = addr_q_scalar * 512 + start_head * 512; - uint64_t q_rope_offset = addr_q_scalar * 64 + start_head * 64; + ctx.q_offset = addr_q_scalar * 512 + start_head * 512; + ctx.q_rope_offset = addr_q_scalar * 64 + start_head * 64; - uint32_t pp_n_scalar = block_size; - uint32_t sub_n_loop = pp_n_scalar / block_size; + ctx.pp_n_scalar = block_size; + ctx.sub_n_loop = ctx.pp_n_scalar / block_size; + ctx.n_loop = (cur_kv_seqlen + ctx.pp_n_scalar - 1) / ctx.pp_n_scalar; - uint32_t n_loop = (cur_kv_seqlen + pp_n_scalar - 1) / pp_n_scalar; + ctx.qk_n = ctx.pp_n_scalar; + ctx.qk_round_n = RoundUp(ctx.qk_n); + ctx.qk_n_2 = ctx.pp_n_scalar; + ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); + ctx.qk_round_n_l1 = RoundUp(ctx.qk_n); + ctx.qk_round_n_2_l1 = RoundUp(ctx.qk_n_2); - uint32_t qk_n = pp_n_scalar; - uint32_t qk_round_n = RoundUp(qk_n); - uint32_t qk_n_2 = pp_n_scalar; - uint32_t qk_round_n_2 = RoundUp(qk_n_2); - uint32_t qk_round_n_l1 = RoundUp(qk_n); - uint32_t qk_round_n_2_l1 = RoundUp(qk_n_2); - uint64_t hidden_size = 576; - if constexpr(tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - hidden_size = 512; + ctx.hidden_size = 576; + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + ctx.hidden_size = 512; } - uint64_t k_round_n = qk_round_n; - uint32_t row_num = cur_head_num * cur_q_seqlen; - m = RoundUp<16>(row_num); + ctx.k_round_n = ctx.qk_round_n; + + ctx.row_num = cur_head_num * cur_q_seqlen; + m = RoundUp<16>(ctx.row_num); + } + + __aicore__ __attribute__((always_inline)) inline void LoadQData(MLAContext &ctx) + { + uint32_t cur_q_seqlen = ctx.cur_q_seqlen; + uint32_t cur_head_num = ctx.cur_head_num; + uint64_t q_offset = ctx.q_offset; + uint64_t q_rope_offset = ctx.q_rope_offset; // copy Q if (cur_q_seqlen == 1) { @@ -747,7 +805,7 @@ class MLAttentionDecoderAic { ) ); } else { - for (uint32_t ii =0; ii < cur_q_seqlen; ii++) { + for (uint32_t ii = 0; ii < cur_q_seqlen; ii++) { AscendC::DataCopy( l1q_buf_addr_tensor[ii * 16], // offset one datablock q_gm_tensor[q_offset + ii * q_heads * 512], @@ -764,7 +822,6 @@ class MLAttentionDecoderAic { ); } } - } if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { gm_to_l1( @@ -795,422 +852,470 @@ class MLAttentionDecoderAic { } SET_FLAG(MTE2, MTE1, EVENT_ID0); WAIT_FLAG(MTE2, MTE1, EVENT_ID0); - for (uint32_t n_idx = 0; n_idx < n_loop + 1; n_idx+=1) { - if (n_idx != n_loop) { - uint32_t l1_kv_pingpong_flag = n_idx % 2; - if (n_idx == (n_loop - 1)) { - qk_n = (cur_kv_seqlen - n_idx * pp_n_scalar); - qk_round_n = RoundUp(qk_n); - qk_round_n_l1 = RoundUp(qk_n); - } - if constexpr(tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - k_round_n = qk_round_n_l1; - } else { - k_round_n = qk_round_n; - } - uint64_t hiddenSize_offset = start_head * cur_q_seqlen * embedding_size; - uint32_t embed_split_size = 128; - uint32_t round_embed_split_size = RoundUp(embed_split_size); + } - /* ************ CUBE1 stage1 ************* */ - - uint32_t block_table_id = (uint32_t)(*(block_tables_gm + - cur_batch * max_num_blocks_per_query + start_kv / block_size + n_idx)); - int64_t kv_offset = (int64_t)block_table_id * block_size * stride_kv; - int64_t kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; - uint32_t q_load_coeff = 1; - q_load_coeff = m; - WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); // wait for v -> L0B - if constexpr(KInputType == InputFormat::ND_FORMAT) { - gm_to_l1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576], - k_gm_tensor[kv_offset], - qk_n, // nValue - qk_round_n, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 512, // dValue - 0, // dstNzMatrixStride, unused - stride_kv // srcDValue - ); - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B - gm_to_l1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576 + 512 * qk_round_n], - k_rope_gm_tensor[kv_offset_rope], - qk_n, // nValue - qk_round_n, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 64, // dValue - 0, // dstNzMatrixStride, unused - stride_kv_rope // srcDValue - ); - } else if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - gm_to_l1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 512], - k_gm_tensor[kv_offset], - qk_round_n_l1, // nValue - block_size, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 512, // dValue - 0, // dstNzMatrixStride, unused - 0 // srcDValue - ); - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B - gm_to_l1( - l1kv_rope_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64], - k_rope_gm_tensor[kv_offset_rope], - qk_round_n, // nValue - block_size, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 64, // dValue - 0, // dstNzMatrixStride, unused - 0 // srcDValue - ); - } else { - gm_to_l1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576], - k_gm_tensor[kv_offset], - qk_round_n, // nValue - block_size, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 512, // dValue - 0, // dstNzMatrixStride, unused - 0 // srcDValue - ); + __aicore__ __attribute__((always_inline)) inline void LoadKVData( + MLAContext &ctx, uint32_t n_idx, uint32_t l1_kv_pingpong_flag) + { + uint32_t qk_n = ctx.qk_n; + uint32_t qk_round_n = ctx.qk_round_n; + uint32_t qk_round_n_l1 = ctx.qk_round_n_l1; + uint32_t cur_batch = ctx.cur_batch; + uint32_t start_kv = ctx.start_kv; + + uint32_t block_table_id = (uint32_t)(*(block_tables_gm + + cur_batch * max_num_blocks_per_query + start_kv / block_size + n_idx)); + int64_t kv_offset = (int64_t)block_table_id * block_size * stride_kv; + int64_t kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; + + WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); // wait for v -> L0B + if constexpr (KInputType == InputFormat::ND_FORMAT) { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576], + k_gm_tensor[kv_offset], + qk_n, // nValue + qk_round_n, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 512, // dValue + 0, // dstNzMatrixStride, unused + stride_kv // srcDValue + ); + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576 + 512 * qk_round_n], + k_rope_gm_tensor[kv_offset_rope], + qk_n, // nValue + qk_round_n, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 64, // dValue + 0, // dstNzMatrixStride, unused + stride_kv_rope // srcDValue + ); + } else if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 512], + k_gm_tensor[kv_offset], + qk_round_n_l1, // nValue + block_size, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 512, // dValue + 0, // dstNzMatrixStride, unused + 0 // srcDValue + ); + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B + gm_to_l1( + l1kv_rope_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64], + k_rope_gm_tensor[kv_offset_rope], + qk_round_n, // nValue + block_size, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 64, // dValue + 0, // dstNzMatrixStride, unused + 0 // srcDValue + ); + } else { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576], + k_gm_tensor[kv_offset], + qk_round_n, // nValue + block_size, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 512, // dValue + 0, // dstNzMatrixStride, unused + 0 // srcDValue + ); - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B - gm_to_l1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576 + 512 * qk_round_n], - k_rope_gm_tensor[kv_offset_rope], - qk_round_n, // nValue - block_size, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 64, // dValue - 0, // dstNzMatrixStride, unused - 0 // srcDValue - ); - } + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576 + 512 * qk_round_n], + k_rope_gm_tensor[kv_offset_rope], + qk_round_n, // nValue + block_size, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 64, // dValue + 0, // dstNzMatrixStride, unused + 0 // srcDValue + ); + } - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); - uint64_t hidden_split_time = (hidden_size + 128 - 1) / 128; - uint64_t embed_split_idx = 0; - for (embed_split_idx = 0; embed_split_idx < hidden_split_time; ++embed_split_idx) { - if (embed_split_idx == 4) { - embed_split_size = 64; - round_embed_split_size = 64; - } - WAIT_FLAG(M, MTE1, embed_split_idx % 2); + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); + } - for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { - l1_to_l0_a( - l0a_buf_tensor[embed_split_idx % 2 * 16384 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], - l1q_buf_addr_tensor[embed_split_idx * m * 128 + loa_load_idx * T_CUBE_MATRIX_SIZE], - 0, - round_embed_split_size / T_BLOCK_SIZE, // repeat - 0, - q_load_coeff / BLOCK_SIZE, // srcStride - 0, - 0 // dstStride - ); - } + __aicore__ __attribute__((always_inline)) inline void ComputeQK( + MLAContext &ctx, uint32_t n_idx, uint32_t l1_kv_pingpong_flag) + { + uint32_t qk_n = ctx.qk_n; + uint32_t qk_round_n = ctx.qk_round_n; + uint32_t qk_round_n_l1 = ctx.qk_round_n_l1; + uint64_t hidden_size = ctx.hidden_size; + uint64_t k_round_n = ctx.k_round_n; + uint32_t row_num = ctx.row_num; + // m is class member, use directly + + uint32_t embed_split_size = 128; + uint32_t round_embed_split_size = RoundUp(embed_split_size); + uint32_t q_load_coeff = m; + + uint64_t hidden_split_time = (hidden_size + 128 - 1) / 128; + uint64_t embed_split_idx = 0; + for (embed_split_idx = 0; embed_split_idx < hidden_split_time; ++embed_split_idx) { + if (embed_split_idx == 4) { + embed_split_size = 64; + round_embed_split_size = 64; + } + WAIT_FLAG(M, MTE1, embed_split_idx % 2); - SET_FLAG(MTE1, M, embed_split_idx % 2); + for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { + l1_to_l0_a( + l0a_buf_tensor[embed_split_idx % 2 * 16384 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], + l1q_buf_addr_tensor[embed_split_idx * m * 128 + loa_load_idx * T_CUBE_MATRIX_SIZE], + 0, + round_embed_split_size / T_BLOCK_SIZE, // repeat + 0, + q_load_coeff / BLOCK_SIZE, // srcStride + 0, + 0 // dstStride + ); + } - if (embed_split_idx == 0) { - WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - } - if (embed_split_idx == 4) { - WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); - } - WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); - l1_to_l0_b( - l0b_buf_tensor[embed_split_idx % 2 * 16384], - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + embed_split_idx * k_round_n * 128], - 0, - round_embed_split_size * k_round_n / T_CUBE_MATRIX_SIZE, // repeat - 0, - 1, // srcStride - 0, - 0 // dstStride - ); - if (embed_split_idx == 4) { - SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); - } - SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); - WAIT_FLAG(MTE1, M, embed_split_idx % 2); - WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); - if (embed_split_idx == 0) { - WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); - } - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - mmad( - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - l0a_buf_tensor[embed_split_idx % 2 * 16384], - l0b_buf_tensor[embed_split_idx % 2 * 16384], - m, // m - qk_round_n_l1, // n - embed_split_size, // k - embed_split_idx == 0 // cmatrixInitVal - ); - } else { - mmad( - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - l0a_buf_tensor[embed_split_idx % 2 * 16384], - l0b_buf_tensor[embed_split_idx % 2 * 16384], - m, // m - qk_n, // n - embed_split_size, // k - embed_split_idx == 0 // cmatrixInitVal - ); - } + SET_FLAG(MTE1, M, embed_split_idx % 2); - PIPE_BARRIER(M); - SET_FLAG(M, MTE1, embed_split_idx % 2); - SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); + if (embed_split_idx == 0) { + WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + } + if (embed_split_idx == 4) { + WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); + } + WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); + l1_to_l0_b( + l0b_buf_tensor[embed_split_idx % 2 * 16384], + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + embed_split_idx * k_round_n * 128], + 0, + round_embed_split_size * k_round_n / T_CUBE_MATRIX_SIZE, // repeat + 0, + 1, // srcStride + 0, + 0 // dstStride + ); + if (embed_split_idx == 4) { + SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); + } + SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); + WAIT_FLAG(MTE1, M, embed_split_idx % 2); + WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); + if (embed_split_idx == 0) { + WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); + } + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + mmad( + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + l0a_buf_tensor[embed_split_idx % 2 * 16384], + l0b_buf_tensor[embed_split_idx % 2 * 16384], + m, // m + qk_round_n_l1, // n + embed_split_size, // k + embed_split_idx == 0 // cmatrixInitVal + ); + } else { + mmad( + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + l0a_buf_tensor[embed_split_idx % 2 * 16384], + l0b_buf_tensor[embed_split_idx % 2 * 16384], + m, // m + qk_n, // n + embed_split_size, // k + embed_split_idx == 0 // cmatrixInitVal + ); + } - // copy S to gm - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - if (embed_split_idx == 3) { - SET_FLAG(M, FIX, l1_kv_pingpong_flag); - WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); + PIPE_BARRIER(M); + SET_FLAG(M, MTE1, embed_split_idx % 2); + SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); - l0c_to_gm( - s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - m, // MSize - qk_n, // NSize - RoundUp<16>(m), // srcStride - qk_round_n // dstStride_dst_D - ); - SET_FLAG(FIX, M, l1_kv_pingpong_flag); - } - } - if (embed_split_idx == 4) { - SET_FLAG(M, FIX, l1_kv_pingpong_flag); - WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); + // copy S to gm + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + if (embed_split_idx == 3) { + SET_FLAG(M, FIX, l1_kv_pingpong_flag); + WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); - l0c_to_gm( - s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - m, // MSize - qk_round_n, // NSize - RoundUp<16>(m), // srcStride - qk_round_n // dstStride_dst_D - ); - SET_FLAG(FIX, M, l1_kv_pingpong_flag); - } + l0c_to_gm( + s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + m, // MSize + qk_n, // NSize + RoundUp<16>(m), // srcStride + qk_round_n // dstStride_dst_D + ); + SET_FLAG(FIX, M, l1_kv_pingpong_flag); } - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - embed_split_idx = 4; - embed_split_size = 64; - round_embed_split_size = 64; - WAIT_FLAG(M, MTE1, embed_split_idx % 2); - - for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { - l1_to_l0_a( - l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], - l1q_rope_buf_addr_tensor[loa_load_idx * CUBE_MATRIX_SIZE], - 0, - round_embed_split_size / BLOCK_SIZE, // repeat - 0, - q_load_coeff / BLOCK_SIZE, // srcStride - 0, - 0 // dstStride - ); - } + } + if (embed_split_idx == 4) { + SET_FLAG(M, FIX, l1_kv_pingpong_flag); + WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); + + l0c_to_gm( + s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + m, // MSize + qk_round_n, // NSize + RoundUp<16>(m), // srcStride + qk_round_n // dstStride_dst_D + ); + SET_FLAG(FIX, M, l1_kv_pingpong_flag); + } + } + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + embed_split_idx = 4; + embed_split_size = 64; + round_embed_split_size = 64; + WAIT_FLAG(M, MTE1, embed_split_idx % 2); + + for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { + l1_to_l0_a( + l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], + l1q_rope_buf_addr_tensor[loa_load_idx * CUBE_MATRIX_SIZE], + 0, + round_embed_split_size / BLOCK_SIZE, // repeat + 0, + q_load_coeff / BLOCK_SIZE, // srcStride + 0, + 0 // dstStride + ); + } - SET_FLAG(MTE1, M, embed_split_idx % 2); + SET_FLAG(MTE1, M, embed_split_idx % 2); - WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); - WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); - l1_to_l0_b( - l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], - l1kv_rope_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64], - 0, - round_embed_split_size * qk_round_n / CUBE_MATRIX_SIZE, // repeat + WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); + WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); + l1_to_l0_b( + l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], + l1kv_rope_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64], + 0, + round_embed_split_size * qk_round_n / CUBE_MATRIX_SIZE, // repeat + 0, + 1, // srcStride + 0, + 0 // dstStride + ); + + SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); + SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); + WAIT_FLAG(MTE1, M, embed_split_idx % 2); + WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); + } + mmad( + mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], + l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], + l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], + m, // m + qk_n, // n + embed_split_size, // k + 1 // cmatrixInitVal + ); + PIPE_BARRIER(M); + SET_FLAG(M, MTE1, embed_split_idx % 2); + SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); + + SET_FLAG(M, FIX, l1_kv_pingpong_flag); + WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + l0c_to_gm( + s_rope_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], + m, // MSize + qk_round_n, // NSize + RoundUp<16>(m), // srcStride + qk_round_n // dstStride_dst_D + ); + } else { + l0c_to_gm( + s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + m, // MSize + qk_round_n, // NSize + RoundUp<16>(m), // srcStride + qk_round_n // dstStride_dst_D + ); + } + SET_FLAG(FIX, M, l1_kv_pingpong_flag); + } + FftsCrossCoreSync(QK_READY_DECODER); + } + + __aicore__ __attribute__((always_inline)) inline void ComputePV( + MLAContext &ctx, uint32_t n_idx) + { + uint32_t qk_n_2 = ctx.qk_n_2; + uint32_t qk_round_n_2 = ctx.qk_round_n_2; + uint32_t qk_round_n_2_l1 = ctx.qk_round_n_2_l1; + uint64_t k_round_n = ctx.k_round_n; + uint32_t row_num = ctx.row_num; + // m is class member, use directly + uint64_t hidden_size = ctx.hidden_size; + + if (n_idx == ctx.n_loop) { + qk_n_2 = (ctx.cur_kv_seqlen - (n_idx - 1) * ctx.pp_n_scalar); + qk_round_n_2 = RoundUp(qk_n_2); + qk_round_n_2_l1 = RoundUp(qk_n_2); + } + k_round_n = qk_round_n_2_l1; + uint32_t l1_kv_pingpong_flag = (n_idx - 1) % 2; + uint32_t l0_p_pingpong_flag = (n_idx - 1) % 2; + uint32_t embed_split_size = 128; + embed_split_loop_v = 4; + uint32_t round_embed_split_size = RoundUp(embed_split_size); + for (uint32_t embed_split_idx = 0; embed_split_idx < embed_split_loop_v; ++embed_split_idx) { + uint32_t l0c_pingpong_flag = (n_idx + embed_split_idx) % 2; + uint32_t l0b_pingpong_flag = (embed_split_idx + 1) % 2; + uint64_t l1kv_offset = embed_split_idx * k_round_n * round_embed_split_size; + WAIT_FLAG(M, MTE1, l0b_pingpong_flag + 2); + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.dstGap = 0; + loadDataParams.startIndex = 0; + loadDataParams.dstFracGap = 0; + if (k_round_n <= round_embed_split_size) { // Nz -> nZ + loadDataParams.repeatTimes = round_embed_split_size / T_BLOCK_SIZE; + loadDataParams.srcStride = k_round_n / T_BLOCK_SIZE; + uint16_t dstGap = sizeof(IN_DTYPE) == 1 ? 1 : 0; + loadDataParams.dstGap = dstGap; + for (uint32_t l0b_load_idx = 0; l0b_load_idx < k_round_n / T_BLOCK_SIZE; ++l0b_load_idx) { + // along embd dim + AscendC::LoadDataWithTranspose( + l0b_buf_tensor[l0b_pingpong_flag * 16384 + l0b_load_idx * RoundUp<16>(embed_split_size) * T_BLOCK_SIZE], + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + l1kv_offset + l0b_load_idx * T_BLOCK_SIZE * T_BLOCK_SIZE], + loadDataParams); + } + } else { + for (uint32_t l0b_load_idx = 0; l0b_load_idx < round_embed_split_size / T_BLOCK_SIZE; ++l0b_load_idx) { + // along kv_len_blk dim + loadDataParams.repeatTimes = qk_round_n_2 / T_BLOCK_SIZE; + loadDataParams.srcStride = 1; + loadDataParams.dstGap = round_embed_split_size / BLOCK_SIZE - 1; + AscendC::LoadDataWithTranspose( + l0b_buf_tensor[l0b_pingpong_flag * 16384 + l0b_load_idx * T_BLOCK_SIZE * T_BLOCK_SIZE], + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + l1kv_offset + l0b_load_idx * qk_round_n_2 * T_BLOCK_SIZE], + loadDataParams); + } + } + if (embed_split_idx == embed_split_loop_v - 1) { + SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); + } + // move p from gm to l1 + uint32_t p_move_head_num = row_num; + if (embed_split_idx == 0) { + WaitFlagDev(SOFTMAX_READY_DECODER); + + WAIT_FLAG(MTE1, MTE2, EVENT_ID7); + gm_to_l1( + l1p_buf_addr_tensor, + p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + ((n_idx - 1) % 2) * TMP_SIZE * T_BLOCK_OFFSET / 2], + p_move_head_num, // nValue + RoundUp(p_move_head_num),// dstNzC0Stride + 0, // dstNzMatrixStride, unused + k_round_n, // dValue + 0, // dstNzMatrixStride, unused + qk_round_n_2 * 2 / sizeof(IN_DTYPE) // srcDValue + ); + SET_FLAG(MTE2, MTE1, EVENT_ID7); + WAIT_FLAG(MTE2, MTE1, EVENT_ID7); + // move p from l1 to l0a + WAIT_FLAG(M, MTE1, l0_p_pingpong_flag); + uint32_t p_load_coeff = RoundUp<16>(p_move_head_num); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + l1_to_l0_a( + l0a_buf_tensor[l0_p_pingpong_flag * 16384], l1p_buf_addr_tensor, RoundUp(p_move_head_num), + qk_round_n_2_l1, // repeat 0, - 1, // srcStride + 0, // srcStride 0, - 0 // dstStride - ); - - SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); - SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); - WAIT_FLAG(MTE1, M, embed_split_idx % 2); - WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); - } - mmad( - mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], - l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], - l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], - m, // m - qk_n, // n - embed_split_size, // k - 1 // cmatrixInitVal + 0 // dstStride ); - PIPE_BARRIER(M); - SET_FLAG(M, MTE1, embed_split_idx % 2); - SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); - - SET_FLAG(M, FIX, l1_kv_pingpong_flag); - WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - l0c_to_gm( - s_rope_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], - m, // MSize - qk_round_n, // NSize - RoundUp<16>(m), // srcStride - qk_round_n // dstStride_dst_D - ); - } else { - l0c_to_gm( - s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - m, // MSize - qk_round_n, // NSize - RoundUp<16>(m), // srcStride - qk_round_n // dstStride_dst_D + } else { + for (uint64_t loa_load_idx = 0; loa_load_idx < p_load_coeff / BLOCK_SIZE; ++loa_load_idx) { + l1_to_l0_a( + l0a_buf_tensor[l0_p_pingpong_flag * 16384 + loa_load_idx * qk_round_n_2 * BLOCK_SIZE], + l1p_buf_addr_tensor[loa_load_idx * T_CUBE_MATRIX_SIZE], + 0, + qk_round_n_2 / T_BLOCK_SIZE, // repeat + 0, + p_load_coeff / BLOCK_SIZE, // srcStride + 0, + 0 // dstStride ); } - SET_FLAG(FIX, M, l1_kv_pingpong_flag); } - FftsCrossCoreSync(QK_READY_DECODER); + SET_FLAG(MTE1, MTE2, EVENT_ID7); } - /* ************ CUBE2 stage1 ************* */ - if (n_idx != 0) { - if (n_idx == n_loop) { - qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); - qk_round_n_2 = RoundUp(qk_n_2); - qk_round_n_2_l1 = RoundUp(qk_n_2); - } - k_round_n = qk_round_n_2_l1; - uint32_t l1_kv_pingpong_flag = (n_idx - 1) % 2; - uint32_t l0_p_pingpong_flag = (n_idx - 1) % 2; - uint32_t embed_split_size = 128; - embed_split_loop_v = 4; - uint32_t round_embed_split_size = RoundUp(embed_split_size); - for (uint32_t embed_split_idx = 0; embed_split_idx < embed_split_loop_v; ++embed_split_idx) { - uint32_t l0c_pingpong_flag = (n_idx + embed_split_idx) % 2; - uint32_t l0b_pingpong_flag = (embed_split_idx + 1) % 2; - uint64_t l1kv_offset = embed_split_idx * k_round_n * round_embed_split_size; - WAIT_FLAG(M, MTE1, l0b_pingpong_flag + 2); - AscendC::LoadData2dTransposeParams loadDataParams; - loadDataParams.dstGap = 0; - loadDataParams.startIndex = 0; - loadDataParams.dstFracGap = 0; - if (k_round_n <= round_embed_split_size) { // Nz -> nZ - loadDataParams.repeatTimes = round_embed_split_size / T_BLOCK_SIZE; - loadDataParams.srcStride = k_round_n / T_BLOCK_SIZE; - uint16_t dstGap = sizeof(IN_DTYPE) == 1 ? 1 : 0; - loadDataParams.dstGap = dstGap; - for (uint32_t l0b_load_idx = 0; l0b_load_idx < k_round_n / T_BLOCK_SIZE; ++l0b_load_idx) { - // along embd dim - AscendC::LoadDataWithTranspose( - l0b_buf_tensor[l0b_pingpong_flag * 16384 + l0b_load_idx * RoundUp<16>(embed_split_size) * T_BLOCK_SIZE], - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + l1kv_offset + l0b_load_idx * T_BLOCK_SIZE * T_BLOCK_SIZE], - loadDataParams); - } - } else { - for (uint32_t l0b_load_idx = 0; l0b_load_idx < round_embed_split_size / T_BLOCK_SIZE; ++l0b_load_idx) { - // along kv_len_blk dim - loadDataParams.repeatTimes = qk_round_n_2 / T_BLOCK_SIZE; - loadDataParams.srcStride = 1; - loadDataParams.dstGap = round_embed_split_size / BLOCK_SIZE - 1; - AscendC::LoadDataWithTranspose( - l0b_buf_tensor[l0b_pingpong_flag * 16384 + l0b_load_idx * T_BLOCK_SIZE * T_BLOCK_SIZE], - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + l1kv_offset + l0b_load_idx * qk_round_n_2 * T_BLOCK_SIZE], - loadDataParams); - } - } - if (embed_split_idx == embed_split_loop_v - 1) { - SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); - } - // move p from gm to l1 - uint32_t p_move_head_num = row_num; - if (embed_split_idx == 0) { - WaitFlagDev(SOFTMAX_READY_DECODER); + SET_FLAG(MTE1, M, l0b_pingpong_flag); + WAIT_FLAG(MTE1, M, l0b_pingpong_flag); + WAIT_FLAG(FIX, M, l0c_pingpong_flag); + mmad( + mm2_l0c_buf_tensor[l0c_pingpong_flag * 16384], + l0a_buf_tensor[l0_p_pingpong_flag * 16384], + l0b_buf_tensor[l0b_pingpong_flag * 16384], + m, // m + embed_split_size, // n + qk_n_2, // k + 1 // cmatrixInitVal + ); + SET_FLAG(M, MTE1, l0b_pingpong_flag + 2); + if (embed_split_idx == embed_split_loop_v - 1) { + SET_FLAG(M, MTE1, l0_p_pingpong_flag); + } + SET_FLAG(M, FIX, l0c_pingpong_flag); + WAIT_FLAG(M, FIX, l0c_pingpong_flag); + + // copy O to gm + l0c_to_gm( + o_tmp_gm_tensor[(uint64_t)block_idx * TMP_SIZE * 2 + embed_split_idx * round_embed_split_size + ((n_idx - 1) % 2) * TMP_SIZE], + mm2_l0c_buf_tensor[l0c_pingpong_flag * 16384], + m, // MSize + RoundUp<16>(embed_split_size), // NSize 32B align + RoundUp<16>(m), // srcStride + round_v // dstStride_dst_D + ); + SET_FLAG(FIX, M, l0c_pingpong_flag); + } + FftsCrossCoreSync(UPDATE_READY_DECODER); + } - WAIT_FLAG(MTE1, MTE2, EVENT_ID7); - gm_to_l1( - l1p_buf_addr_tensor, - p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + ((n_idx - 1) % 2) * TMP_SIZE * T_BLOCK_OFFSET / 2], - p_move_head_num, // nValue - RoundUp(p_move_head_num),// dstNzC0Stride - 0, // dstNzMatrixStride, unused - k_round_n, // dValue - 0, // dstNzMatrixStride, unused - qk_round_n_2 * 2 / sizeof(IN_DTYPE) // srcDValue - ); - SET_FLAG(MTE2, MTE1, EVENT_ID7); - WAIT_FLAG(MTE2, MTE1, EVENT_ID7); - // move p from l1 to l0a - WAIT_FLAG(M, MTE1, l0_p_pingpong_flag); - uint32_t p_load_coeff = RoundUp<16>(p_move_head_num); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - l1_to_l0_a( - l0a_buf_tensor[l0_p_pingpong_flag * 16384], l1p_buf_addr_tensor, RoundUp(p_move_head_num), - qk_round_n_2_l1, // repeat - 0, - 0, // srcStride - 0, - 0 // dstStride - ); - } else { - for (uint64_t loa_load_idx = 0; loa_load_idx < p_load_coeff / BLOCK_SIZE; ++loa_load_idx) { - l1_to_l0_a( - l0a_buf_tensor[l0_p_pingpong_flag * 16384 + loa_load_idx * qk_round_n_2 * BLOCK_SIZE], - l1p_buf_addr_tensor[loa_load_idx * T_CUBE_MATRIX_SIZE], - 0, - qk_round_n_2 / T_BLOCK_SIZE, // repeat - 0, - p_load_coeff / BLOCK_SIZE, // srcStride - 0, - 0 // dstStride - ); - } - } - SET_FLAG(MTE1, MTE2, EVENT_ID7); - } - SET_FLAG(MTE1, M, l0b_pingpong_flag); - WAIT_FLAG(MTE1, M, l0b_pingpong_flag); - WAIT_FLAG(FIX, M, l0c_pingpong_flag); - mmad( - mm2_l0c_buf_tensor[l0c_pingpong_flag * 16384], - l0a_buf_tensor[l0_p_pingpong_flag * 16384], - l0b_buf_tensor[l0b_pingpong_flag * 16384], - m, // m - embed_split_size, // n - qk_n_2, // k - 1 // cmatrixInitVal - ); - SET_FLAG(M, MTE1, l0b_pingpong_flag + 2); - if (embed_split_idx == embed_split_loop_v - 1) { - SET_FLAG(M, MTE1, l0_p_pingpong_flag); - } - SET_FLAG(M, FIX, l0c_pingpong_flag); - WAIT_FLAG(M, FIX, l0c_pingpong_flag); + __aicore__ __attribute__((always_inline)) inline void InnerRunCubeMLA(uint32_t cur_batch, uint32_t start_head, uint32_t cur_head_num, + uint32_t start_kv, uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t offset_tiling) + { + MLAContext ctx; + InitMLAContext(ctx, cur_batch, start_head, cur_head_num, + start_kv, cur_q_seqlen, cur_kv_seqlen, offset_tiling); - // copy O to gm - l0c_to_gm( - o_tmp_gm_tensor[(uint64_t)block_idx * TMP_SIZE * 2 + embed_split_idx * round_embed_split_size + ((n_idx - 1) % 2) * TMP_SIZE], - mm2_l0c_buf_tensor[l0c_pingpong_flag * 16384], - m, // MSize - RoundUp<16>(embed_split_size), // NSize 32B align - RoundUp<16>(m), // srcStride - round_v // dstStride_dst_D - ); - SET_FLAG(FIX, M, l0c_pingpong_flag); + LoadQData(ctx); + for (uint32_t n_idx = 0; n_idx < ctx.n_loop + 1; n_idx += 1) { + if (n_idx != ctx.n_loop) { + uint32_t l1_kv_pingpong_flag = n_idx % 2; + if (n_idx == (ctx.n_loop - 1)) { + ctx.qk_n = (ctx.cur_kv_seqlen - n_idx * ctx.pp_n_scalar); + ctx.qk_round_n = RoundUp(ctx.qk_n); + ctx.qk_round_n_l1 = RoundUp(ctx.qk_n); } - FftsCrossCoreSync(UPDATE_READY_DECODER); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + ctx.k_round_n = ctx.qk_round_n_l1; + } else { + ctx.k_round_n = ctx.qk_round_n; + } + + LoadKVData(ctx, n_idx, l1_kv_pingpong_flag); + ComputeQK(ctx, n_idx, l1_kv_pingpong_flag); + } + if (n_idx != 0) { + ComputePV(ctx, n_idx); } } } + __aicore__ __attribute__((always_inline)) inline void InnerRunCubeMLATP1(uint32_t cur_batch, uint32_t start_head, uint32_t cur_head_num, uint32_t start_kv, uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t offset_tiling) { From e091e26b19624a6cd7a86ed1a4c9a964bd313663 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Fri, 7 Aug 2026 15:07:53 +0800 Subject: [PATCH 08/24] feat: you hua mla code --- .../op_kernel/multi_latent_attention.h | 314 +++++++++--------- 1 file changed, 152 insertions(+), 162 deletions(-) diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h index 7f9bb38..8d8aa36 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h @@ -769,91 +769,175 @@ class MLAttentionDecoderAic { m = RoundUp<16>(ctx.row_num); } - __aicore__ __attribute__((always_inline)) inline void LoadQData(MLAContext &ctx) + // 公共函数：将 Q 主体数据从 GM 搬运到 L1（ND→NZ 格式转换） + // 覆盖三种场景： + // 1. cur_q_seqlen == 1 → gm_to_l1 单矩阵搬运 + // 2. cur_q_seqlen > 1 && q_heads < 128 → DataCopy 批量多矩阵搬运 + // 3. cur_q_seqlen > 1 && q_heads >= 128 → for循环逐token搬运（规避 stride 位宽限制） + __aicore__ __attribute__((always_inline)) inline void LoadQMainFromGMToL1( + AscendC::LocalTensor &l1_dst, + AscendC::GlobalTensor &gm_src, + uint64_t src_offset, + uint32_t cur_q_seqlen, + uint32_t cur_head_num) { - uint32_t cur_q_seqlen = ctx.cur_q_seqlen; - uint32_t cur_head_num = ctx.cur_head_num; - uint64_t q_offset = ctx.q_offset; - uint64_t q_rope_offset = ctx.q_rope_offset; - - // copy Q if (cur_q_seqlen == 1) { gm_to_l1( - l1q_buf_addr_tensor, - q_gm_tensor[q_offset], - cur_head_num, // nValue - RoundUp<16>(cur_head_num),// dstNzC0Stride - 0, // dstNzMatrixStride, unused - 512, // dValue - 0, // dstNzMatrixStride, unused - 512 // srcDValue + l1_dst, + gm_src[src_offset], + cur_head_num, // nValue + RoundUp<16>(cur_head_num), // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 512, // dValue + 0, // dstNzMatrixStride, unused + 512 // srcDValue ); } else { if (q_heads < 128) { AscendC::DataCopy( - l1q_buf_addr_tensor, - q_gm_tensor[q_offset], + l1_dst, + gm_src[src_offset], AscendC::Nd2NzParams( - cur_q_seqlen, // ndNum - cur_head_num, // nValue - 512, // dValue - 512 * q_heads, // srcNdMatrixStride - 512, // srcDValue + cur_q_seqlen, // ndNum + cur_head_num, // nValue + 512, // dValue + 512 * q_heads, // srcNdMatrixStride + 512, // srcDValue RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride - cur_q_seqlen, // dstNzNStride - 16 // dstNzMatrixStride + cur_q_seqlen, // dstNzNStride + 16 // dstNzMatrixStride ) ); } else { for (uint32_t ii = 0; ii < cur_q_seqlen; ii++) { AscendC::DataCopy( - l1q_buf_addr_tensor[ii * 16], // offset one datablock - q_gm_tensor[q_offset + ii * q_heads * 512], + l1_dst[ii * 16], // offset one datablock + gm_src[src_offset + ii * q_heads * 512], AscendC::Nd2NzParams( - 1, // ndNum - cur_head_num, // nValue - 512, // dValue - 0, // srcNdMatrixStride - 512, // srcDValue + 1, // ndNum + cur_head_num, // nValue + 512, // dValue + 0, // srcNdMatrixStride + 512, // srcDValue RoundUp<16>(cur_q_seqlen * cur_head_num), // dstNzC0Stride - cur_q_seqlen, // dstNzNStride - 16 // dstNzMatrixStride + cur_q_seqlen, // dstNzNStride + 16 // dstNzMatrixStride ) ); } } } + } + + // 公共函数：将 Q Rope 数据从 GM 搬运到 L1（ND→NZ 格式转换） + // INT8 场景：用 gm_to_l1 搬到独立的 l1q_rope_buf_addr_tensor + // 非INT8 场景：用 DataCopy 搬到 l1q_buf_addr_tensor 的 Q 主体之后 + __aicore__ __attribute__((always_inline)) inline void LoadQRopeFromGMToL1( + AscendC::LocalTensor &l1_q, + AscendC::LocalTensor &l1_q_rope, + AscendC::GlobalTensor &gm_src, + uint64_t src_offset, + uint32_t cur_q_seqlen, + uint32_t cur_head_num) + { if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { gm_to_l1( - l1q_rope_buf_addr_tensor, - q_rope_gm_tensor[q_rope_offset], - cur_head_num, // nValue - RoundUp<16>(cur_head_num),// dstNzC0Stride - 0, // dstNzMatrixStride, unused - 64, // dValue - 0, // dstNzMatrixStride, unused - 64 // srcDValue + l1_q_rope, + gm_src[src_offset], + cur_head_num, // nValue + RoundUp<16>(cur_head_num), // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 64, // dValue + 0, // dstNzMatrixStride, unused + 64 // srcDValue ); } else { AscendC::DataCopy( - l1q_buf_addr_tensor[RoundUp<16>(cur_head_num * cur_q_seqlen) * 512], - q_rope_gm_tensor[q_rope_offset], + l1_q[RoundUp<16>(cur_head_num * cur_q_seqlen) * 512], + gm_src[src_offset], AscendC::Nd2NzParams( - cur_head_num, // ndNum, 32 - cur_q_seqlen, // nValue, 4 - 64, // dValue - 64, // srcNdMatrixStride - 64 * q_heads, // srcDValue - RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride - 1, // dstNzNStride - 16 * cur_q_seqlen // dstNzMatrixStride + cur_head_num, // ndNum + cur_q_seqlen, // nValue + 64, // dValue + 64, // srcNdMatrixStride + 64 * q_heads, // srcDValue + RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride + 1, // dstNzNStride + 16 * cur_q_seqlen // dstNzMatrixStride ) ); } + } + + __aicore__ __attribute__((always_inline)) inline void LoadQData(MLAContext &ctx) + { + uint32_t cur_q_seqlen = ctx.cur_q_seqlen; + uint32_t cur_head_num = ctx.cur_head_num; + uint64_t q_offset = ctx.q_offset; + uint64_t q_rope_offset = ctx.q_rope_offset; + + // copy Q + LoadQMainFromGMToL1(l1q_buf_addr_tensor, q_gm_tensor, q_offset, cur_q_seqlen, cur_head_num); + LoadQRopeFromGMToL1(l1q_buf_addr_tensor, l1q_rope_buf_addr_tensor, q_rope_gm_tensor, q_rope_offset, cur_q_seqlen, cur_head_num); SET_FLAG(MTE2, MTE1, EVENT_ID0); WAIT_FLAG(MTE2, MTE1, EVENT_ID0); } + // 公共函数：将 KV 主体数据从 GM 搬运到 L1 + // ND_FORMAT 场景：ND→NZ 格式转换（gm_to_l1） + // INT8 / 非INT8 NZ 场景：NZ→NZ 格式搬运（gm_to_l1） + template + __aicore__ __attribute__((always_inline)) inline void LoadKVMainFromGMToL1( + AscendC::LocalTensor &l1_dst, + AscendC::GlobalTensor &gm_src, + uint32_t n_value, // 实际行数 + uint32_t dst_nz_c0_stride, // L1 对齐行数（dstNzC0Stride） + uint32_t d_value, // 列数（512） + uint32_t src_d_value, // GM 行 stride（ND场景为stride_kv，NZ场景为0） + bool is_nd_to_nz) // true=ND→NZ, false=NZ→NZ + { + if (is_nd_to_nz) { + gm_to_l1( + l1_dst, + gm_src, + n_value, // nValue + dst_nz_c0_stride, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + d_value, // dValue + 0, // dstNzMatrixStride, unused + src_d_value // srcDValue + ); + } else { + gm_to_l1( + l1_dst, + gm_src, + n_value, // nValue + dst_nz_c0_stride, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + d_value, // dValue + 0, // dstNzMatrixStride, unused + 0 // srcDValue (NZ→NZ always 0) + ); + } + } + + // 公共函数：将 KV Rope 数据从 GM 搬运到 L1 + // ND_FORMAT 场景：ND→NZ 格式转换（gm_to_l1），d_value=64, src_d_value=stride_kv_rope + // INT8 / 非INT8 NZ 场景：NZ→NZ 格式搬运（gm_to_l1），d_value=64, src_d_value=0 + template + __aicore__ __attribute__((always_inline)) inline void LoadKVRopeFromGMToL1( + AscendC::LocalTensor &l1_dst, + AscendC::GlobalTensor &gm_src, + uint32_t n_value, // 实际行数 + uint32_t dst_nz_c0_stride, // L1 对齐行数（dstNzC0Stride） + uint32_t d_value, // 列数（64） + uint32_t src_d_value, // GM 行 stride（ND场景为stride_kv_rope，NZ场景为0） + bool is_nd_to_nz) // true=ND→NZ, false=NZ→NZ + { + // 逻辑与 LoadKVMainFromGMToL1 完全一致，仅参数不同（d_value=64） + LoadKVMainFromGMToL1(l1_dst, gm_src, n_value, dst_nz_c0_stride, d_value, src_d_value, is_nd_to_nz); + } + __aicore__ __attribute__((always_inline)) inline void LoadKVData( MLAContext &ctx, uint32_t n_idx, uint32_t l1_kv_pingpong_flag) { @@ -870,75 +954,41 @@ class MLAttentionDecoderAic { WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); // wait for v -> L0B if constexpr (KInputType == InputFormat::ND_FORMAT) { - gm_to_l1( + // 分支1: ND→NZ，K 主体搬到 l1kv_buf，Rope 紧随其后 + LoadKVMainFromGMToL1( l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576], k_gm_tensor[kv_offset], - qk_n, // nValue - qk_round_n, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 512, // dValue - 0, // dstNzMatrixStride, unused - stride_kv // srcDValue - ); + qk_n, qk_round_n, 512, stride_kv, true /*ND→NZ*/); SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B - gm_to_l1( + LoadKVRopeFromGMToL1( l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576 + 512 * qk_round_n], k_rope_gm_tensor[kv_offset_rope], - qk_n, // nValue - qk_round_n, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 64, // dValue - 0, // dstNzMatrixStride, unused - stride_kv_rope // srcDValue - ); + qk_n, qk_round_n, 64, stride_kv_rope, true /*ND→NZ*/); } else if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - gm_to_l1( + // 分支2: NZ→NZ，K 主体搬到 l1kv_buf(512布局)，Rope 搬到独立 l1kv_rope_buf + LoadKVMainFromGMToL1( l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 512], k_gm_tensor[kv_offset], - qk_round_n_l1, // nValue - block_size, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 512, // dValue - 0, // dstNzMatrixStride, unused - 0 // srcDValue - ); + qk_round_n_l1, block_size, 512, 0, false /*NZ→NZ*/); SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B - gm_to_l1( + LoadKVRopeFromGMToL1( l1kv_rope_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64], k_rope_gm_tensor[kv_offset_rope], - qk_round_n, // nValue - block_size, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 64, // dValue - 0, // dstNzMatrixStride, unused - 0 // srcDValue - ); + qk_round_n, block_size, 64, 0, false /*NZ→NZ*/); } else { - gm_to_l1( + // 分支3: NZ→NZ，K 主体搬到 l1kv_buf(576布局)，Rope 紧随其后 + LoadKVMainFromGMToL1( l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576], k_gm_tensor[kv_offset], - qk_round_n, // nValue - block_size, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 512, // dValue - 0, // dstNzMatrixStride, unused - 0 // srcDValue - ); - + qk_round_n, block_size, 512, 0, false /*NZ→NZ*/); SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B - gm_to_l1( + LoadKVRopeFromGMToL1( l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576 + 512 * qk_round_n], k_rope_gm_tensor[kv_offset_rope], - qk_round_n, // nValue - block_size, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 64, // dValue - 0, // dstNzMatrixStride, unused - 0 // srcDValue - ); + qk_round_n, block_size, 64, 0, false /*NZ→NZ*/); } SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); @@ -1340,68 +1390,8 @@ class MLAttentionDecoderAic { m = RoundUp<16>(row_num); // copy Q - if (cur_q_seqlen == 1) { - gm_to_l1( - l1q_buf_addr_tensor, - q_gm_tensor[q_offset], - cur_head_num, // nValue - RoundUp<16>(cur_head_num),// dstNzC0Stride - 0, // dstNzMatrixStride, unused - 512, // dValue - 0, // dstNzMatrixStride, unused - 512 // srcDValue - ); - } else { - if (q_heads < 128) { - AscendC::DataCopy( - l1q_buf_addr_tensor, - q_gm_tensor[q_offset], - AscendC::Nd2NzParams( - cur_q_seqlen, // ndNum - cur_head_num, // nValue - 512, // dValue - 512 * q_heads, // srcNdMatrixStride - 512, // srcDValue - RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride - cur_q_seqlen, // dstNzNStride - 16 // dstNzMatrixStride - ) - ); - } else { - for (uint32_t ii =0; ii < cur_q_seqlen; ii++) { - AscendC::DataCopy( - l1q_buf_addr_tensor[ii * 16], // offset one datablock - q_gm_tensor[q_offset + ii * q_heads * 512], - AscendC::Nd2NzParams( - 1, // ndNum - cur_head_num, // nValue - 512, // dValue - 0, // srcNdMatrixStride - 512, // srcDValue - RoundUp<16>(cur_q_seqlen * cur_head_num), // dstNzC0Stride - cur_q_seqlen, // dstNzNStride - 16 // dstNzMatrixStride - ) - ); - } - } - - } - - AscendC::DataCopy( - l1q_buf_addr_tensor[RoundUp<16>(cur_head_num * cur_q_seqlen) * 512], - q_rope_gm_tensor[q_rope_offset], - AscendC::Nd2NzParams( - cur_head_num, // ndNum, 32 - cur_q_seqlen, // nValue, 4 - 64, // dValue - 64, // srcNdMatrixStride - 64 * q_heads, // srcDValue - RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride - 1, // dstNzNStride - 16 * cur_q_seqlen // dstNzMatrixStride - ) - ); + LoadQMainFromGMToL1(l1q_buf_addr_tensor, q_gm_tensor, q_offset, cur_q_seqlen, cur_head_num); + LoadQRopeFromGMToL1(l1q_buf_addr_tensor, l1q_rope_buf_addr_tensor, q_rope_gm_tensor, q_rope_offset, cur_q_seqlen, cur_head_num); SET_FLAG(MTE2, MTE1, EVENT_ID0); WAIT_FLAG(MTE2, MTE1, EVENT_ID0); From f30393329b4892d126a3d0be6e8eeb66126115c8 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Mon, 10 Aug 2026 14:21:30 +0800 Subject: [PATCH 09/24] feat: modify the mla --- .../op_kernel/multi_latent_attention.h | 562 +++--------------- .../op_kernel/multi_latent_attention_arch32.h | 445 ++++++++++++++ .../op_kernel/multi_latent_attention_bs.h | 348 +++++++++++ 3 files changed, 873 insertions(+), 482 deletions(-) create mode 100644 xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_arch32.h create mode 100644 xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h index 8d8aa36..0c9e62f 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h @@ -769,106 +769,11 @@ class MLAttentionDecoderAic { m = RoundUp<16>(ctx.row_num); } - // 公共函数：将 Q 主体数据从 GM 搬运到 L1（ND→NZ 格式转换） - // 覆盖三种场景： - // 1. cur_q_seqlen == 1 → gm_to_l1 单矩阵搬运 - // 2. cur_q_seqlen > 1 && q_heads < 128 → DataCopy 批量多矩阵搬运 - // 3. cur_q_seqlen > 1 && q_heads >= 128 → for循环逐token搬运（规避 stride 位宽限制） - __aicore__ __attribute__((always_inline)) inline void LoadQMainFromGMToL1( - AscendC::LocalTensor &l1_dst, - AscendC::GlobalTensor &gm_src, - uint64_t src_offset, - uint32_t cur_q_seqlen, - uint32_t cur_head_num) - { - if (cur_q_seqlen == 1) { - gm_to_l1( - l1_dst, - gm_src[src_offset], - cur_head_num, // nValue - RoundUp<16>(cur_head_num), // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 512, // dValue - 0, // dstNzMatrixStride, unused - 512 // srcDValue - ); - } else { - if (q_heads < 128) { - AscendC::DataCopy( - l1_dst, - gm_src[src_offset], - AscendC::Nd2NzParams( - cur_q_seqlen, // ndNum - cur_head_num, // nValue - 512, // dValue - 512 * q_heads, // srcNdMatrixStride - 512, // srcDValue - RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride - cur_q_seqlen, // dstNzNStride - 16 // dstNzMatrixStride - ) - ); - } else { - for (uint32_t ii = 0; ii < cur_q_seqlen; ii++) { - AscendC::DataCopy( - l1_dst[ii * 16], // offset one datablock - gm_src[src_offset + ii * q_heads * 512], - AscendC::Nd2NzParams( - 1, // ndNum - cur_head_num, // nValue - 512, // dValue - 0, // srcNdMatrixStride - 512, // srcDValue - RoundUp<16>(cur_q_seqlen * cur_head_num), // dstNzC0Stride - cur_q_seqlen, // dstNzNStride - 16 // dstNzMatrixStride - ) - ); - } - } - } - } - - // 公共函数：将 Q Rope 数据从 GM 搬运到 L1（ND→NZ 格式转换） - // INT8 场景：用 gm_to_l1 搬到独立的 l1q_rope_buf_addr_tensor - // 非INT8 场景：用 DataCopy 搬到 l1q_buf_addr_tensor 的 Q 主体之后 - __aicore__ __attribute__((always_inline)) inline void LoadQRopeFromGMToL1( - AscendC::LocalTensor &l1_q, - AscendC::LocalTensor &l1_q_rope, - AscendC::GlobalTensor &gm_src, - uint64_t src_offset, - uint32_t cur_q_seqlen, - uint32_t cur_head_num) - { - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - gm_to_l1( - l1_q_rope, - gm_src[src_offset], - cur_head_num, // nValue - RoundUp<16>(cur_head_num), // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 64, // dValue - 0, // dstNzMatrixStride, unused - 64 // srcDValue - ); - } else { - AscendC::DataCopy( - l1_q[RoundUp<16>(cur_head_num * cur_q_seqlen) * 512], - gm_src[src_offset], - AscendC::Nd2NzParams( - cur_head_num, // ndNum - cur_q_seqlen, // nValue - 64, // dValue - 64, // srcNdMatrixStride - 64 * q_heads, // srcDValue - RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride - 1, // dstNzNStride - 16 * cur_q_seqlen // dstNzMatrixStride - ) - ); - } - } + // === 第一层：LoadQData 编排层 === + // 业务函数（LoadQMainFromGMToL1 / LoadQRopeFromGMToL1 / PlatformSetQLoadComplete） + // 已迁移至 bs.h / arch32.h，通过 include 展开 + // 业务函数：加载 Q 数据（主体 + Rope）从 GM 到 L1 __aicore__ __attribute__((always_inline)) inline void LoadQData(MLAContext &ctx) { uint32_t cur_q_seqlen = ctx.cur_q_seqlen; @@ -879,64 +784,11 @@ class MLAttentionDecoderAic { // copy Q LoadQMainFromGMToL1(l1q_buf_addr_tensor, q_gm_tensor, q_offset, cur_q_seqlen, cur_head_num); LoadQRopeFromGMToL1(l1q_buf_addr_tensor, l1q_rope_buf_addr_tensor, q_rope_gm_tensor, q_rope_offset, cur_q_seqlen, cur_head_num); - SET_FLAG(MTE2, MTE1, EVENT_ID0); - WAIT_FLAG(MTE2, MTE1, EVENT_ID0); + PlatformSetQLoadComplete(); } - // 公共函数：将 KV 主体数据从 GM 搬运到 L1 - // ND_FORMAT 场景：ND→NZ 格式转换（gm_to_l1） - // INT8 / 非INT8 NZ 场景：NZ→NZ 格式搬运（gm_to_l1） - template - __aicore__ __attribute__((always_inline)) inline void LoadKVMainFromGMToL1( - AscendC::LocalTensor &l1_dst, - AscendC::GlobalTensor &gm_src, - uint32_t n_value, // 实际行数 - uint32_t dst_nz_c0_stride, // L1 对齐行数（dstNzC0Stride） - uint32_t d_value, // 列数（512） - uint32_t src_d_value, // GM 行 stride（ND场景为stride_kv，NZ场景为0） - bool is_nd_to_nz) // true=ND→NZ, false=NZ→NZ - { - if (is_nd_to_nz) { - gm_to_l1( - l1_dst, - gm_src, - n_value, // nValue - dst_nz_c0_stride, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - d_value, // dValue - 0, // dstNzMatrixStride, unused - src_d_value // srcDValue - ); - } else { - gm_to_l1( - l1_dst, - gm_src, - n_value, // nValue - dst_nz_c0_stride, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - d_value, // dValue - 0, // dstNzMatrixStride, unused - 0 // srcDValue (NZ→NZ always 0) - ); - } - } - - // 公共函数：将 KV Rope 数据从 GM 搬运到 L1 - // ND_FORMAT 场景：ND→NZ 格式转换（gm_to_l1），d_value=64, src_d_value=stride_kv_rope - // INT8 / 非INT8 NZ 场景：NZ→NZ 格式搬运（gm_to_l1），d_value=64, src_d_value=0 - template - __aicore__ __attribute__((always_inline)) inline void LoadKVRopeFromGMToL1( - AscendC::LocalTensor &l1_dst, - AscendC::GlobalTensor &gm_src, - uint32_t n_value, // 实际行数 - uint32_t dst_nz_c0_stride, // L1 对齐行数（dstNzC0Stride） - uint32_t d_value, // 列数（64） - uint32_t src_d_value, // GM 行 stride（ND场景为stride_kv_rope，NZ场景为0） - bool is_nd_to_nz) // true=ND→NZ, false=NZ→NZ - { - // 逻辑与 LoadKVMainFromGMToL1 完全一致，仅参数不同（d_value=64） - LoadKVMainFromGMToL1(l1_dst, gm_src, n_value, dst_nz_c0_stride, d_value, src_d_value, is_nd_to_nz); - } + // === 第一层：LoadKVData 编排层 === + // 业务函数已迁移至 bs.h / arch32.h __aicore__ __attribute__((always_inline)) inline void LoadKVData( MLAContext &ctx, uint32_t n_idx, uint32_t l1_kv_pingpong_flag) @@ -952,15 +804,14 @@ class MLAttentionDecoderAic { int64_t kv_offset = (int64_t)block_table_id * block_size * stride_kv; int64_t kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; - WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); // wait for v -> L0B + PlatformWaitKVLoadReady(l1_kv_pingpong_flag); if constexpr (KInputType == InputFormat::ND_FORMAT) { // 分支1: ND→NZ，K 主体搬到 l1kv_buf，Rope 紧随其后 LoadKVMainFromGMToL1( l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576], k_gm_tensor[kv_offset], qk_n, qk_round_n, 512, stride_kv, true /*ND→NZ*/); - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B + PlatformSetKVMainLoadComplete(l1_kv_pingpong_flag); LoadKVRopeFromGMToL1( l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576 + 512 * qk_round_n], k_rope_gm_tensor[kv_offset_rope], @@ -971,8 +822,7 @@ class MLAttentionDecoderAic { l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 512], k_gm_tensor[kv_offset], qk_round_n_l1, block_size, 512, 0, false /*NZ→NZ*/); - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B + PlatformSetKVMainLoadComplete(l1_kv_pingpong_flag); LoadKVRopeFromGMToL1( l1kv_rope_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64], k_rope_gm_tensor[kv_offset_rope], @@ -983,353 +833,101 @@ class MLAttentionDecoderAic { l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576], k_gm_tensor[kv_offset], qk_round_n, block_size, 512, 0, false /*NZ→NZ*/); - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B + PlatformSetKVMainLoadComplete(l1_kv_pingpong_flag); LoadKVRopeFromGMToL1( l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576 + 512 * qk_round_n], k_rope_gm_tensor[kv_offset_rope], qk_round_n, block_size, 64, 0, false /*NZ→NZ*/); } - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); + PlatformSetKVRopeLoadComplete(l1_kv_pingpong_flag); } + // ==================== ComputeQK 三层架构拆分 ==================== + // === 第一层：QKParams 结构体 === + struct QKParams { + uint32_t qk_n; + uint32_t qk_round_n; + uint32_t qk_round_n_l1; + uint64_t hidden_size; + uint64_t k_round_n; + uint32_t row_num; + uint32_t embed_split_size; + uint32_t round_embed_split_size; + uint32_t q_load_coeff; + uint64_t hidden_split_time; + uint32_t l1_kv_pingpong_flag; + }; + + // [QK 第三层平台函数 + 第二层业务函数] 已迁移至 arch32.h / bs.h + // （include 指令统一放在 PVParams 结构体之后，确保所有 Params 类型可见） + + // === 第一层：ComputeQK 编排层 === __aicore__ __attribute__((always_inline)) inline void ComputeQK( MLAContext &ctx, uint32_t n_idx, uint32_t l1_kv_pingpong_flag) { - uint32_t qk_n = ctx.qk_n; - uint32_t qk_round_n = ctx.qk_round_n; - uint32_t qk_round_n_l1 = ctx.qk_round_n_l1; - uint64_t hidden_size = ctx.hidden_size; - uint64_t k_round_n = ctx.k_round_n; - uint32_t row_num = ctx.row_num; - // m is class member, use directly - - uint32_t embed_split_size = 128; - uint32_t round_embed_split_size = RoundUp(embed_split_size); - uint32_t q_load_coeff = m; - - uint64_t hidden_split_time = (hidden_size + 128 - 1) / 128; - uint64_t embed_split_idx = 0; - for (embed_split_idx = 0; embed_split_idx < hidden_split_time; ++embed_split_idx) { - if (embed_split_idx == 4) { - embed_split_size = 64; - round_embed_split_size = 64; - } - WAIT_FLAG(M, MTE1, embed_split_idx % 2); - - for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { - l1_to_l0_a( - l0a_buf_tensor[embed_split_idx % 2 * 16384 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], - l1q_buf_addr_tensor[embed_split_idx * m * 128 + loa_load_idx * T_CUBE_MATRIX_SIZE], - 0, - round_embed_split_size / T_BLOCK_SIZE, // repeat - 0, - q_load_coeff / BLOCK_SIZE, // srcStride - 0, - 0 // dstStride - ); - } - - SET_FLAG(MTE1, M, embed_split_idx % 2); - - if (embed_split_idx == 0) { - WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - } - if (embed_split_idx == 4) { - WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); - } - WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); - l1_to_l0_b( - l0b_buf_tensor[embed_split_idx % 2 * 16384], - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + embed_split_idx * k_round_n * 128], - 0, - round_embed_split_size * k_round_n / T_CUBE_MATRIX_SIZE, // repeat - 0, - 1, // srcStride - 0, - 0 // dstStride - ); - if (embed_split_idx == 4) { - SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); - } - SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); - WAIT_FLAG(MTE1, M, embed_split_idx % 2); - WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); - if (embed_split_idx == 0) { - WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); - } - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - mmad( - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - l0a_buf_tensor[embed_split_idx % 2 * 16384], - l0b_buf_tensor[embed_split_idx % 2 * 16384], - m, // m - qk_round_n_l1, // n - embed_split_size, // k - embed_split_idx == 0 // cmatrixInitVal - ); - } else { - mmad( - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - l0a_buf_tensor[embed_split_idx % 2 * 16384], - l0b_buf_tensor[embed_split_idx % 2 * 16384], - m, // m - qk_n, // n - embed_split_size, // k - embed_split_idx == 0 // cmatrixInitVal - ); - } + QKParams params; + InitQKParams(ctx, l1_kv_pingpong_flag, params); - PIPE_BARRIER(M); - SET_FLAG(M, MTE1, embed_split_idx % 2); - SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); - - // copy S to gm - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - if (embed_split_idx == 3) { - SET_FLAG(M, FIX, l1_kv_pingpong_flag); - WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); - - l0c_to_gm( - s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - m, // MSize - qk_n, // NSize - RoundUp<16>(m), // srcStride - qk_round_n // dstStride_dst_D - ); - SET_FLAG(FIX, M, l1_kv_pingpong_flag); - } - } + for (uint64_t embed_split_idx = 0; embed_split_idx < params.hidden_split_time; ++embed_split_idx) { if (embed_split_idx == 4) { - SET_FLAG(M, FIX, l1_kv_pingpong_flag); - WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); - - l0c_to_gm( - s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - m, // MSize - qk_round_n, // NSize - RoundUp<16>(m), // srcStride - qk_round_n // dstStride_dst_D - ); - SET_FLAG(FIX, M, l1_kv_pingpong_flag); + params.embed_split_size = 64; + params.round_embed_split_size = 64; } + LoadQDataToL0A(params, embed_split_idx, false); + LoadKVDataToL0B(params, embed_split_idx, false); + ComputeQKMMad(params, embed_split_idx, false); + CopyQKResultToGM(params, embed_split_idx, n_idx, false); } - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - embed_split_idx = 4; - embed_split_size = 64; - round_embed_split_size = 64; - WAIT_FLAG(M, MTE1, embed_split_idx % 2); - - for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { - l1_to_l0_a( - l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], - l1q_rope_buf_addr_tensor[loa_load_idx * CUBE_MATRIX_SIZE], - 0, - round_embed_split_size / BLOCK_SIZE, // repeat - 0, - q_load_coeff / BLOCK_SIZE, // srcStride - 0, - 0 // dstStride - ); - } - SET_FLAG(MTE1, M, embed_split_idx % 2); - - WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); - WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); - l1_to_l0_b( - l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], - l1kv_rope_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64], - 0, - round_embed_split_size * qk_round_n / CUBE_MATRIX_SIZE, // repeat - 0, - 1, // srcStride - 0, - 0 // dstStride - ); - - SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); - SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); - WAIT_FLAG(MTE1, M, embed_split_idx % 2); - WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); - } - mmad( - mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], - l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], - l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], - m, // m - qk_n, // n - embed_split_size, // k - 1 // cmatrixInitVal - ); - PIPE_BARRIER(M); - SET_FLAG(M, MTE1, embed_split_idx % 2); - SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); - - SET_FLAG(M, FIX, l1_kv_pingpong_flag); - WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - l0c_to_gm( - s_rope_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], - m, // MSize - qk_round_n, // NSize - RoundUp<16>(m), // srcStride - qk_round_n // dstStride_dst_D - ); - } else { - l0c_to_gm( - s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - m, // MSize - qk_round_n, // NSize - RoundUp<16>(m), // srcStride - qk_round_n // dstStride_dst_D - ); - } - SET_FLAG(FIX, M, l1_kv_pingpong_flag); - } + ComputeQRope(params, n_idx); FftsCrossCoreSync(QK_READY_DECODER); } + // ==================== ComputePV 三层架构拆分 ==================== + // === 第一层：PVParams 结构体 === + struct PVParams { + uint32_t qk_n_2; + uint32_t qk_round_n_2; + uint32_t qk_round_n_2_l1; + uint64_t k_round_n; + uint32_t row_num; + uint64_t hidden_size; + uint32_t l1_kv_pingpong_flag; + uint32_t l0_p_pingpong_flag; + uint32_t embed_split_size; + uint32_t round_embed_split_size; + uint32_t l0c_pingpong_flag; + uint32_t l0b_pingpong_flag; + uint64_t l1kv_offset; + }; + + // [PV 第三层平台函数 + 第二层业务函数] 已迁移至 arch32.h / bs.h + // === 以下 include 第三层和第二层子文件（放在 QKParams/PVParams 之后，确保类型可见）=== + #include "multi_latent_attention_arch32.h" + #include "multi_latent_attention_bs.h" + + // === 第一层：ComputePV 编排层 === __aicore__ __attribute__((always_inline)) inline void ComputePV( MLAContext &ctx, uint32_t n_idx) { - uint32_t qk_n_2 = ctx.qk_n_2; - uint32_t qk_round_n_2 = ctx.qk_round_n_2; - uint32_t qk_round_n_2_l1 = ctx.qk_round_n_2_l1; - uint64_t k_round_n = ctx.k_round_n; - uint32_t row_num = ctx.row_num; - // m is class member, use directly - uint64_t hidden_size = ctx.hidden_size; - - if (n_idx == ctx.n_loop) { - qk_n_2 = (ctx.cur_kv_seqlen - (n_idx - 1) * ctx.pp_n_scalar); - qk_round_n_2 = RoundUp(qk_n_2); - qk_round_n_2_l1 = RoundUp(qk_n_2); - } - k_round_n = qk_round_n_2_l1; - uint32_t l1_kv_pingpong_flag = (n_idx - 1) % 2; - uint32_t l0_p_pingpong_flag = (n_idx - 1) % 2; - uint32_t embed_split_size = 128; - embed_split_loop_v = 4; - uint32_t round_embed_split_size = RoundUp(embed_split_size); + PVParams params; + InitPVParams(ctx, n_idx, params); + for (uint32_t embed_split_idx = 0; embed_split_idx < embed_split_loop_v; ++embed_split_idx) { - uint32_t l0c_pingpong_flag = (n_idx + embed_split_idx) % 2; - uint32_t l0b_pingpong_flag = (embed_split_idx + 1) % 2; - uint64_t l1kv_offset = embed_split_idx * k_round_n * round_embed_split_size; - WAIT_FLAG(M, MTE1, l0b_pingpong_flag + 2); - AscendC::LoadData2dTransposeParams loadDataParams; - loadDataParams.dstGap = 0; - loadDataParams.startIndex = 0; - loadDataParams.dstFracGap = 0; - if (k_round_n <= round_embed_split_size) { // Nz -> nZ - loadDataParams.repeatTimes = round_embed_split_size / T_BLOCK_SIZE; - loadDataParams.srcStride = k_round_n / T_BLOCK_SIZE; - uint16_t dstGap = sizeof(IN_DTYPE) == 1 ? 1 : 0; - loadDataParams.dstGap = dstGap; - for (uint32_t l0b_load_idx = 0; l0b_load_idx < k_round_n / T_BLOCK_SIZE; ++l0b_load_idx) { - // along embd dim - AscendC::LoadDataWithTranspose( - l0b_buf_tensor[l0b_pingpong_flag * 16384 + l0b_load_idx * RoundUp<16>(embed_split_size) * T_BLOCK_SIZE], - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + l1kv_offset + l0b_load_idx * T_BLOCK_SIZE * T_BLOCK_SIZE], - loadDataParams); - } - } else { - for (uint32_t l0b_load_idx = 0; l0b_load_idx < round_embed_split_size / T_BLOCK_SIZE; ++l0b_load_idx) { - // along kv_len_blk dim - loadDataParams.repeatTimes = qk_round_n_2 / T_BLOCK_SIZE; - loadDataParams.srcStride = 1; - loadDataParams.dstGap = round_embed_split_size / BLOCK_SIZE - 1; - AscendC::LoadDataWithTranspose( - l0b_buf_tensor[l0b_pingpong_flag * 16384 + l0b_load_idx * T_BLOCK_SIZE * T_BLOCK_SIZE], - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + l1kv_offset + l0b_load_idx * qk_round_n_2 * T_BLOCK_SIZE], - loadDataParams); - } - } - if (embed_split_idx == embed_split_loop_v - 1) { - SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); - } - // move p from gm to l1 - uint32_t p_move_head_num = row_num; + params.l0c_pingpong_flag = (n_idx + embed_split_idx) % 2; + params.l0b_pingpong_flag = (embed_split_idx + 1) % 2; + params.l1kv_offset = embed_split_idx * params.k_round_n * params.round_embed_split_size; + + LoadKVTransposeToL0B(params, embed_split_idx); + if (embed_split_idx == 0) { - WaitFlagDev(SOFTMAX_READY_DECODER); - - WAIT_FLAG(MTE1, MTE2, EVENT_ID7); - gm_to_l1( - l1p_buf_addr_tensor, - p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + ((n_idx - 1) % 2) * TMP_SIZE * T_BLOCK_OFFSET / 2], - p_move_head_num, // nValue - RoundUp(p_move_head_num),// dstNzC0Stride - 0, // dstNzMatrixStride, unused - k_round_n, // dValue - 0, // dstNzMatrixStride, unused - qk_round_n_2 * 2 / sizeof(IN_DTYPE) // srcDValue - ); - SET_FLAG(MTE2, MTE1, EVENT_ID7); - WAIT_FLAG(MTE2, MTE1, EVENT_ID7); - // move p from l1 to l0a - WAIT_FLAG(M, MTE1, l0_p_pingpong_flag); - uint32_t p_load_coeff = RoundUp<16>(p_move_head_num); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - l1_to_l0_a( - l0a_buf_tensor[l0_p_pingpong_flag * 16384], l1p_buf_addr_tensor, RoundUp(p_move_head_num), - qk_round_n_2_l1, // repeat - 0, - 0, // srcStride - 0, - 0 // dstStride - ); - } else { - for (uint64_t loa_load_idx = 0; loa_load_idx < p_load_coeff / BLOCK_SIZE; ++loa_load_idx) { - l1_to_l0_a( - l0a_buf_tensor[l0_p_pingpong_flag * 16384 + loa_load_idx * qk_round_n_2 * BLOCK_SIZE], - l1p_buf_addr_tensor[loa_load_idx * T_CUBE_MATRIX_SIZE], - 0, - qk_round_n_2 / T_BLOCK_SIZE, // repeat - 0, - p_load_coeff / BLOCK_SIZE, // srcStride - 0, - 0 // dstStride - ); - } - } - SET_FLAG(MTE1, MTE2, EVENT_ID7); - } - SET_FLAG(MTE1, M, l0b_pingpong_flag); - WAIT_FLAG(MTE1, M, l0b_pingpong_flag); - WAIT_FLAG(FIX, M, l0c_pingpong_flag); - mmad( - mm2_l0c_buf_tensor[l0c_pingpong_flag * 16384], - l0a_buf_tensor[l0_p_pingpong_flag * 16384], - l0b_buf_tensor[l0b_pingpong_flag * 16384], - m, // m - embed_split_size, // n - qk_n_2, // k - 1 // cmatrixInitVal - ); - SET_FLAG(M, MTE1, l0b_pingpong_flag + 2); - if (embed_split_idx == embed_split_loop_v - 1) { - SET_FLAG(M, MTE1, l0_p_pingpong_flag); + LoadPDataToL0A(params, n_idx); } - SET_FLAG(M, FIX, l0c_pingpong_flag); - WAIT_FLAG(M, FIX, l0c_pingpong_flag); - - // copy O to gm - l0c_to_gm( - o_tmp_gm_tensor[(uint64_t)block_idx * TMP_SIZE * 2 + embed_split_idx * round_embed_split_size + ((n_idx - 1) % 2) * TMP_SIZE], - mm2_l0c_buf_tensor[l0c_pingpong_flag * 16384], - m, // MSize - RoundUp<16>(embed_split_size), // NSize 32B align - RoundUp<16>(m), // srcStride - round_v // dstStride_dst_D - ); - SET_FLAG(FIX, M, l0c_pingpong_flag); + + ComputePVMmad(params, embed_split_idx); + + CopyPVResultToGM(params, embed_split_idx, n_idx); } FftsCrossCoreSync(UPDATE_READY_DECODER); } diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_arch32.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_arch32.h new file mode 100644 index 0000000..a669a5e --- /dev/null +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_arch32.h @@ -0,0 +1,445 @@ +// Copyright 2025 The xLLM Authors. 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 +// +// https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE +// +// 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. +// ============================================================================== +// +// 文件：multi_latent_attention_arch32.h +// 层级：第三层 - 平台处理函数（Platform* 系列） +// 说明：本文件被 include 到 MLAttentionDecoderAic 类的 private 区域内， +// 包含所有与 Ascend C 平台 API 直接交互的函数（gm_to_l1 / l1_to_l0_a / +// l1_to_l0_b / l0c_to_gm / mmad / SET_FLAG / WAIT_FLAG / PIPE_BARRIER 等）。 +// 这些函数不包含业务逻辑，仅负责平台相关的数据搬运与同步操作。 + +// === 第三层：平台处理函数 === + +// 平台函数：设置 Q 数据搬运完成同步（MTE2→MTE1） +__aicore__ __attribute__((always_inline)) inline void PlatformSetQLoadComplete() +{ + SET_FLAG(MTE2, MTE1, EVENT_ID0); + WAIT_FLAG(MTE2, MTE1, EVENT_ID0); +} + +// 平台函数：等待 KV 数据搬运就绪（MTE1→MTE2，等待 V→L0B 完成） +__aicore__ __attribute__((always_inline)) inline void PlatformWaitKVLoadReady(uint32_t l1_kv_pingpong_flag) +{ + WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); +} + +// 平台函数：设置 KV 主体搬运完成 + 等待 V→L0B 完成（MTE2→MTE1 set, MTE1→MTE2 wait） +__aicore__ __attribute__((always_inline)) inline void PlatformSetKVMainLoadComplete(uint32_t l1_kv_pingpong_flag) +{ + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); +} + +// 平台函数：设置 KV Rope 搬运完成（MTE2→MTE1） +__aicore__ __attribute__((always_inline)) inline void PlatformSetKVRopeLoadComplete(uint32_t l1_kv_pingpong_flag) +{ + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); +} + +// 平台函数：Q 数据从 L1 加载到 L0A +__aicore__ __attribute__((always_inline)) inline void PlatformLoadQToL0A( + uint32_t embed_split_idx, uint32_t round_embed_split_size, + uint32_t q_load_coeff, bool is_rope) +{ + uint32_t loa_repeat = is_rope ? round_embed_split_size / BLOCK_SIZE + : round_embed_split_size / T_BLOCK_SIZE; + for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { + if (is_rope) { + l1_to_l0_a( + l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], + l1q_rope_buf_addr_tensor[loa_load_idx * CUBE_MATRIX_SIZE], + 0, loa_repeat, 0, q_load_coeff / BLOCK_SIZE, 0, 0); + } else { + l1_to_l0_a( + l0a_buf_tensor[embed_split_idx % 2 * 16384 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], + l1q_buf_addr_tensor[embed_split_idx * q_load_coeff * 128 + loa_load_idx * T_CUBE_MATRIX_SIZE], + 0, loa_repeat, 0, q_load_coeff / BLOCK_SIZE, 0, 0); + } + } +} + +// 平台函数：设置 Q 加载完成（MTE1→M） +__aicore__ __attribute__((always_inline)) inline void PlatformSetQLoadComplete( + uint32_t embed_split_idx) +{ + SET_FLAG(MTE1, M, embed_split_idx % 2); +} + +// 平台函数：等待 Q 加载完成（M→MTE1） +__aicore__ __attribute__((always_inline)) inline void PlatformWaitQLoadComplete( + uint32_t embed_split_idx) +{ + WAIT_FLAG(M, MTE1, embed_split_idx % 2); +} + +// 平台函数：KV 数据从 L1 加载到 L0B +__aicore__ __attribute__((always_inline)) inline void PlatformLoadKVToL0B( + uint32_t embed_split_idx, uint32_t round_embed_split_size, + uint32_t k_round_n, uint32_t qk_round_n, + uint32_t l1_kv_pingpong_flag, uint64_t hidden_size, bool is_rope) +{ + uint32_t l0b_repeat; + uint32_t l1_src_offset; + if (is_rope) { + l0b_repeat = round_embed_split_size * qk_round_n / CUBE_MATRIX_SIZE; + l1_src_offset = l1_kv_pingpong_flag * 128 * 64; + } else { + l0b_repeat = round_embed_split_size * k_round_n / T_CUBE_MATRIX_SIZE; + l1_src_offset = l1_kv_pingpong_flag * 128 * hidden_size + embed_split_idx * k_round_n * 128; + } + + if (embed_split_idx == 0) { + WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + } + if (embed_split_idx == 4) { + WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); + } + WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); + + if (is_rope) { + l1_to_l0_b( + l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], + l1kv_rope_buf_addr_tensor[l1_src_offset], + 0, l0b_repeat, 0, 1, 0, 0); + } else { + l1_to_l0_b( + l0b_buf_tensor[embed_split_idx % 2 * 16384], + l1kv_buf_addr_tensor[l1_src_offset], + 0, l0b_repeat, 0, 1, 0, 0); + } + + if (embed_split_idx == 4) { + SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); + } + SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); +} + +// 平台函数：QK MMAD 计算（含同步） +__aicore__ __attribute__((always_inline)) inline void PlatformComputeQKMMad( + uint32_t embed_split_idx, uint32_t embed_split_size, + uint32_t m_value, uint32_t qk_n, uint32_t qk_round_n_l1, + uint32_t l1_kv_pingpong_flag, bool is_rope) +{ + WAIT_FLAG(MTE1, M, embed_split_idx % 2); + WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); + if (embed_split_idx == 0) { + WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); + } + + if (is_rope) { + mmad( + mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], + l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], + l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], + m_value, qk_n, embed_split_size, 1); + } else { + uint32_t n_value = qk_n; + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + n_value = qk_round_n_l1; + } + mmad( + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + l0a_buf_tensor[embed_split_idx % 2 * 16384], + l0b_buf_tensor[embed_split_idx % 2 * 16384], + m_value, n_value, embed_split_size, embed_split_idx == 0); + } + + PIPE_BARRIER(M); + SET_FLAG(M, MTE1, embed_split_idx % 2); + SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); +} + +// 平台函数：L0C 结果拷贝到 GM（含同步） +__aicore__ __attribute__((always_inline)) inline void PlatformCopyQKResultToGM( + uint32_t embed_split_idx, uint32_t m_value, uint32_t qk_n, + uint32_t qk_round_n, uint32_t l1_kv_pingpong_flag, + uint32_t n_idx, bool is_rope) +{ + if (is_rope) { + SET_FLAG(M, FIX, l1_kv_pingpong_flag); + WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + l0c_to_gm( + s_rope_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], + m_value, qk_round_n, RoundUp<16>(m_value), qk_round_n); + } else { + l0c_to_gm( + s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + m_value, qk_round_n, RoundUp<16>(m_value), qk_round_n); + } + SET_FLAG(FIX, M, l1_kv_pingpong_flag); + return; + } + + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + if (embed_split_idx == 3) { + SET_FLAG(M, FIX, l1_kv_pingpong_flag); + WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); + l0c_to_gm( + s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + m_value, qk_n, RoundUp<16>(m_value), qk_round_n); + SET_FLAG(FIX, M, l1_kv_pingpong_flag); + } + } + if (embed_split_idx == 4) { + SET_FLAG(M, FIX, l1_kv_pingpong_flag); + WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); + l0c_to_gm( + s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + m_value, qk_round_n, RoundUp<16>(m_value), qk_round_n); + SET_FLAG(FIX, M, l1_kv_pingpong_flag); + } +} + +// 平台函数：KV L1→L0B 转置加载（k_round_n <= round_embed_split_size 分支） +__aicore__ __attribute__((always_inline)) inline void PlatformLoadKVTransposeSmallN( + uint32_t l0b_pingpong_flag, uint32_t l1_kv_pingpong_flag, + uint64_t l1kv_offset, uint64_t k_round_n, uint64_t hidden_size, + uint32_t embed_split_size, uint32_t round_embed_split_size) +{ + WAIT_FLAG(M, MTE1, l0b_pingpong_flag + 2); + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.dstGap = 0; + loadDataParams.startIndex = 0; + loadDataParams.dstFracGap = 0; + loadDataParams.repeatTimes = round_embed_split_size / T_BLOCK_SIZE; + loadDataParams.srcStride = k_round_n / T_BLOCK_SIZE; + uint16_t dstGap = sizeof(IN_DTYPE) == 1 ? 1 : 0; + loadDataParams.dstGap = dstGap; + for (uint32_t l0b_load_idx = 0; l0b_load_idx < k_round_n / T_BLOCK_SIZE; ++l0b_load_idx) { + AscendC::LoadDataWithTranspose( + l0b_buf_tensor[l0b_pingpong_flag * 16384 + l0b_load_idx * RoundUp<16>(embed_split_size) * T_BLOCK_SIZE], + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + l1kv_offset + l0b_load_idx * T_BLOCK_SIZE * T_BLOCK_SIZE], + loadDataParams); + } +} + +// 平台函数：KV L1→L0B 转置加载（k_round_n > round_embed_split_size 分支） +__aicore__ __attribute__((always_inline)) inline void PlatformLoadKVTransposeLargeN( + uint32_t l0b_pingpong_flag, uint32_t l1_kv_pingpong_flag, + uint64_t l1kv_offset, uint64_t hidden_size, + uint32_t qk_round_n_2, uint32_t round_embed_split_size) +{ + WAIT_FLAG(M, MTE1, l0b_pingpong_flag + 2); + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.dstGap = 0; + loadDataParams.startIndex = 0; + loadDataParams.dstFracGap = 0; + for (uint32_t l0b_load_idx = 0; l0b_load_idx < round_embed_split_size / T_BLOCK_SIZE; ++l0b_load_idx) { + loadDataParams.repeatTimes = qk_round_n_2 / T_BLOCK_SIZE; + loadDataParams.srcStride = 1; + loadDataParams.dstGap = round_embed_split_size / BLOCK_SIZE - 1; + AscendC::LoadDataWithTranspose( + l0b_buf_tensor[l0b_pingpong_flag * 16384 + l0b_load_idx * T_BLOCK_SIZE * T_BLOCK_SIZE], + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + l1kv_offset + l0b_load_idx * qk_round_n_2 * T_BLOCK_SIZE], + loadDataParams); + } +} + +// 平台函数：KV 加载完成后设置 MTE1→MTE2 同步 +__aicore__ __attribute__((always_inline)) inline void PlatformSetKVLoadComplete( + bool is_last_split, uint32_t l1_kv_pingpong_flag) +{ + if (is_last_split) { + SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); + } +} + +// 平台函数：P 数据从 GM 加载到 L1 +__aicore__ __attribute__((always_inline)) inline void PlatformLoadPFromGMToL1( + uint32_t n_idx, uint32_t row_num, uint64_t k_round_n, + uint32_t qk_round_n_2) +{ + WaitFlagDev(SOFTMAX_READY_DECODER); + WAIT_FLAG(MTE1, MTE2, EVENT_ID7); + gm_to_l1( + l1p_buf_addr_tensor, + p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + ((n_idx - 1) % 2) * TMP_SIZE * T_BLOCK_OFFSET / 2], + row_num, + RoundUp(row_num), + 0, + k_round_n, + 0, + qk_round_n_2 * 2 / sizeof(IN_DTYPE)); + SET_FLAG(MTE2, MTE1, EVENT_ID7); + WAIT_FLAG(MTE2, MTE1, EVENT_ID7); +} + +// 平台函数：P 数据从 L1 加载到 L0A（INT8 路径） +__aicore__ __attribute__((always_inline)) inline void PlatformLoadPToL0AInt8( + uint32_t l0_p_pingpong_flag, uint32_t row_num, + uint32_t qk_round_n_2_l1) +{ + WAIT_FLAG(M, MTE1, l0_p_pingpong_flag); + l1_to_l0_a( + l0a_buf_tensor[l0_p_pingpong_flag * 16384], l1p_buf_addr_tensor, + RoundUp(row_num), qk_round_n_2_l1, 0, 0, 0, 0); +} + +// 平台函数：P 数据从 L1 加载到 L0A（非 INT8 路径） +__aicore__ __attribute__((always_inline)) inline void PlatformLoadPToL0AGeneral( + uint32_t l0_p_pingpong_flag, uint32_t row_num, + uint32_t qk_round_n_2) +{ + WAIT_FLAG(M, MTE1, l0_p_pingpong_flag); + uint32_t p_load_coeff = RoundUp<16>(row_num); + for (uint64_t loa_load_idx = 0; loa_load_idx < p_load_coeff / BLOCK_SIZE; ++loa_load_idx) { + l1_to_l0_a( + l0a_buf_tensor[l0_p_pingpong_flag * 16384 + loa_load_idx * qk_round_n_2 * BLOCK_SIZE], + l1p_buf_addr_tensor[loa_load_idx * T_CUBE_MATRIX_SIZE], + 0, qk_round_n_2 / T_BLOCK_SIZE, 0, p_load_coeff / BLOCK_SIZE, 0, 0); + } +} + +// 平台函数：P 加载完成后设置 MTE1→MTE2 同步 +__aicore__ __attribute__((always_inline)) inline void PlatformSetPLoadComplete() +{ + SET_FLAG(MTE1, MTE2, EVENT_ID7); +} + +// 平台函数：PV MMAD 计算（含同步） +__aicore__ __attribute__((always_inline)) inline void PlatformComputePVMmad( + uint32_t l0b_pingpong_flag, uint32_t l0c_pingpong_flag, + uint32_t l0_p_pingpong_flag, uint32_t embed_split_size, + uint32_t qk_n_2, bool is_last_split) +{ + SET_FLAG(MTE1, M, l0b_pingpong_flag); + WAIT_FLAG(MTE1, M, l0b_pingpong_flag); + WAIT_FLAG(FIX, M, l0c_pingpong_flag); + mmad( + mm2_l0c_buf_tensor[l0c_pingpong_flag * 16384], + l0a_buf_tensor[l0_p_pingpong_flag * 16384], + l0b_buf_tensor[l0b_pingpong_flag * 16384], + m, embed_split_size, qk_n_2, 1); + SET_FLAG(M, MTE1, l0b_pingpong_flag + 2); + if (is_last_split) { + SET_FLAG(M, MTE1, l0_p_pingpong_flag); + } +} + +// 平台函数：L0C 结果拷贝到 GM（含同步） +__aicore__ __attribute__((always_inline)) inline void PlatformCopyPVResultToGM( + uint32_t l0c_pingpong_flag, uint32_t embed_split_idx, + uint32_t n_idx, uint32_t round_embed_split_size) +{ + SET_FLAG(M, FIX, l0c_pingpong_flag); + WAIT_FLAG(M, FIX, l0c_pingpong_flag); + l0c_to_gm( + o_tmp_gm_tensor[(uint64_t)block_idx * TMP_SIZE * 2 + embed_split_idx * round_embed_split_size + ((n_idx - 1) % 2) * TMP_SIZE], + mm2_l0c_buf_tensor[l0c_pingpong_flag * 16384], + m, RoundUp<16>(round_embed_split_size), RoundUp<16>(m), round_v); + SET_FLAG(FIX, M, l0c_pingpong_flag); +} + +// ==================== Rope 平台函数（INT8 专有）==================== + +// 平台函数：Q_Rope 从 L1 加载到 L0A（for 循环 l1_to_l0_a） +__aicore__ __attribute__((always_inline)) inline void PlatformLoadQRopeToL0A( + uint32_t embed_split_idx, uint32_t q_load_coeff, + uint32_t round_embed_split_size) +{ + WAIT_FLAG(M, MTE1, embed_split_idx % 2); + for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { + l1_to_l0_a( + l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], + l1q_rope_buf_addr_tensor[loa_load_idx * CUBE_MATRIX_SIZE], + 0, + round_embed_split_size / BLOCK_SIZE, + 0, + q_load_coeff / BLOCK_SIZE, + 0, + 0 + ); + } + SET_FLAG(MTE1, M, embed_split_idx % 2); +} + +// 平台函数：KV_Rope 从 L1 加载到 L0B（单次 l1_to_l0_b + 同步） +__aicore__ __attribute__((always_inline)) inline void PlatformLoadKVRopeToL0B( + uint32_t embed_split_idx, uint32_t l1_kv_pingpong_flag, + uint32_t round_embed_split_size, uint32_t qk_round_n) +{ + WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); + WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); + l1_to_l0_b( + l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], + l1kv_rope_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64], + 0, + round_embed_split_size * qk_round_n / CUBE_MATRIX_SIZE, + 0, + 1, + 0, + 0 + ); + SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); + SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); +} + +// 平台函数：Rope MMAD 计算（含同步） +__aicore__ __attribute__((always_inline)) inline void PlatformComputeQRopeMMad( + uint32_t embed_split_idx, uint32_t l1_kv_pingpong_flag, + uint32_t m_value, uint32_t qk_n, uint32_t embed_split_size) +{ + WAIT_FLAG(MTE1, M, embed_split_idx % 2); + WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); + } + mmad( + mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], + l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], + l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], + m_value, + qk_n, + embed_split_size, + 1 + ); + PIPE_BARRIER(M); + SET_FLAG(M, MTE1, embed_split_idx % 2); + SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); +} + +// 平台函数：Rope 结果 L0C 拷贝到 GM（含同步） +__aicore__ __attribute__((always_inline)) inline void PlatformCopyQRopeResultToGM( + uint32_t l1_kv_pingpong_flag, uint32_t n_idx, + uint32_t m_value, uint32_t qk_round_n) +{ + SET_FLAG(M, FIX, l1_kv_pingpong_flag); + WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + l0c_to_gm( + s_rope_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], + m_value, + qk_round_n, + RoundUp<16>(m_value), + qk_round_n + ); + } else { + l0c_to_gm( + s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + m_value, + qk_round_n, + RoundUp<16>(m_value), + qk_round_n + ); + } + SET_FLAG(FIX, M, l1_kv_pingpong_flag); +} \ No newline at end of file diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h new file mode 100644 index 0000000..8395eba --- /dev/null +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h @@ -0,0 +1,348 @@ +// Copyright 2025 The xLLM Authors. 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 +// +// https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE +// +// 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. +// ============================================================================== +// +// 文件：multi_latent_attention_bs.h +// 层级：第二层 - 业务逻辑函数（Load*/Init*/Compute*/Copy* 系列） +// 说明：本文件被 include 到 MLAttentionDecoderAic 类的 private 区域内， +// 包含所有纯流程控制函数，这些函数负责调用第三层 Platform* 系列函数 +// 完成具体的业务编排，不直接调用平台 API。 +// 依赖：第三层函数（Platform* 系列）来自 multi_latent_attention_arch32.h + +// === 第二层：业务逻辑函数 === + +// 业务函数：将 Q 主体数据从 GM 搬运到 L1（ND→NZ 格式转换） +// 覆盖三种场景： +// 1. cur_q_seqlen == 1 → gm_to_l1 单矩阵搬运 +// 2. cur_q_seqlen > 1 && q_heads < 128 → DataCopy 批量多矩阵搬运 +// 3. cur_q_seqlen > 1 && q_heads >= 128 → for循环逐token搬运（规避 stride 位宽限制） +__aicore__ __attribute__((always_inline)) inline void LoadQMainFromGMToL1( + AscendC::LocalTensor &l1_dst, + AscendC::GlobalTensor &gm_src, + uint64_t src_offset, + uint32_t cur_q_seqlen, + uint32_t cur_head_num) +{ + if (cur_q_seqlen == 1) { + gm_to_l1( + l1_dst, + gm_src[src_offset], + cur_head_num, // nValue + RoundUp<16>(cur_head_num), // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 512, // dValue + 0, // dstNzMatrixStride, unused + 512 // srcDValue + ); + } else { + if (q_heads < 128) { + AscendC::DataCopy( + l1_dst, + gm_src[src_offset], + AscendC::Nd2NzParams( + cur_q_seqlen, // ndNum + cur_head_num, // nValue + 512, // dValue + 512 * q_heads, // srcNdMatrixStride + 512, // srcDValue + RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride + cur_q_seqlen, // dstNzNStride + 16 // dstNzMatrixStride + ) + ); + } else { + for (uint32_t ii = 0; ii < cur_q_seqlen; ii++) { + AscendC::DataCopy( + l1_dst[ii * 16], // offset one datablock + gm_src[src_offset + ii * q_heads * 512], + AscendC::Nd2NzParams( + 1, // ndNum + cur_head_num, // nValue + 512, // dValue + 0, // srcNdMatrixStride + 512, // srcDValue + RoundUp<16>(cur_q_seqlen * cur_head_num), // dstNzC0Stride + cur_q_seqlen, // dstNzNStride + 16 // dstNzMatrixStride + ) + ); + } + } + } +} + +// 业务函数：将 Q Rope 数据从 GM 搬运到 L1（ND→NZ 格式转换） +// INT8 场景：用 gm_to_l1 搬到独立的 l1q_rope_buf_addr_tensor +// 非INT8 场景：用 DataCopy 搬到 l1q_buf_addr_tensor 的 Q 主体之后 +__aicore__ __attribute__((always_inline)) inline void LoadQRopeFromGMToL1( + AscendC::LocalTensor &l1_q, + AscendC::LocalTensor &l1_q_rope, + AscendC::GlobalTensor &gm_src, + uint64_t src_offset, + uint32_t cur_q_seqlen, + uint32_t cur_head_num) +{ + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + gm_to_l1( + l1_q_rope, + gm_src[src_offset], + cur_head_num, // nValue + RoundUp<16>(cur_head_num), // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 64, // dValue + 0, // dstNzMatrixStride, unused + 64 // srcDValue + ); + } else { + AscendC::DataCopy( + l1_q[RoundUp<16>(cur_head_num * cur_q_seqlen) * 512], + gm_src[src_offset], + AscendC::Nd2NzParams( + cur_head_num, // ndNum + cur_q_seqlen, // nValue + 64, // dValue + 64, // srcNdMatrixStride + 64 * q_heads, // srcDValue + RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride + 1, // dstNzNStride + 16 * cur_q_seqlen // dstNzMatrixStride + ) + ); + } +} + +// 业务函数：将 KV 主体数据从 GM 搬运到 L1 +// ND_FORMAT 场景：ND→NZ 格式转换（gm_to_l1） +// INT8 / 非INT8 NZ 场景：NZ→NZ 格式搬运（gm_to_l1） +template +__aicore__ __attribute__((always_inline)) inline void LoadKVMainFromGMToL1( + AscendC::LocalTensor l1_dst, + AscendC::GlobalTensor gm_src, + uint32_t n_value, // 实际行数 + uint32_t dst_nz_c0_stride, // L1 对齐行数（dstNzC0Stride） + uint32_t d_value, // 列数（512） + uint32_t src_d_value, // GM 行 stride（ND场景为stride_kv，NZ场景为0） + bool is_nd_to_nz) // true=ND→NZ, false=NZ→NZ +{ + if (is_nd_to_nz) { + gm_to_l1( + l1_dst, + gm_src, + n_value, // nValue + dst_nz_c0_stride, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + d_value, // dValue + 0, // dstNzMatrixStride, unused + src_d_value // srcDValue + ); + } else { + gm_to_l1( + l1_dst, + gm_src, + n_value, // nValue + dst_nz_c0_stride, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + d_value, // dValue + 0, // dstNzMatrixStride, unused + 0 // srcDValue (NZ→NZ always 0) + ); + } +} + +// 业务函数：将 KV Rope 数据从 GM 搬运到 L1 +// ND_FORMAT 场景：ND→NZ 格式转换（gm_to_l1），d_value=64, src_d_value=stride_kv_rope +// INT8 / 非INT8 NZ 场景：NZ→NZ 格式搬运（gm_to_l1），d_value=64, src_d_value=0 +template +__aicore__ __attribute__((always_inline)) inline void LoadKVRopeFromGMToL1( + AscendC::LocalTensor l1_dst, + AscendC::GlobalTensor gm_src, + uint32_t n_value, // 实际行数 + uint32_t dst_nz_c0_stride, // L1 对齐行数（dstNzC0Stride） + uint32_t d_value, // 列数（64） + uint32_t src_d_value, // GM 行 stride（ND场景为stride_kv_rope，NZ场景为0） + bool is_nd_to_nz) // true=ND→NZ, false=NZ→NZ +{ + // 逻辑与 LoadKVMainFromGMToL1 完全一致，仅参数不同（d_value=64） + LoadKVMainFromGMToL1(l1_dst, gm_src, n_value, dst_nz_c0_stride, d_value, src_d_value, is_nd_to_nz); +} + +// 业务函数：初始化 QK 参数 +__aicore__ __attribute__((always_inline)) inline void InitQKParams( + MLAContext &ctx, uint32_t l1_kv_pingpong_flag, QKParams ¶ms) +{ + params.qk_n = ctx.qk_n; + params.qk_round_n = ctx.qk_round_n; + params.qk_round_n_l1 = ctx.qk_round_n_l1; + params.hidden_size = ctx.hidden_size; + params.k_round_n = ctx.k_round_n; + params.row_num = ctx.row_num; + params.embed_split_size = 128; + params.round_embed_split_size = RoundUp(params.embed_split_size); + params.q_load_coeff = m; + params.hidden_split_time = (params.hidden_size + 128 - 1) / 128; + params.l1_kv_pingpong_flag = l1_kv_pingpong_flag; +} + +// 业务函数：Q 数据从 L1 加载到 L0A +__aicore__ __attribute__((always_inline)) inline void LoadQDataToL0A( + const QKParams ¶ms, uint32_t embed_split_idx, bool is_rope) +{ + PlatformWaitQLoadComplete(embed_split_idx); + PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, + params.q_load_coeff, is_rope); + PlatformSetQLoadComplete(embed_split_idx); +} + +// 业务函数：KV 数据从 L1 加载到 L0B +__aicore__ __attribute__((always_inline)) inline void LoadKVDataToL0B( + const QKParams ¶ms, uint32_t embed_split_idx, bool is_rope) +{ + PlatformLoadKVToL0B(embed_split_idx, params.round_embed_split_size, + params.k_round_n, params.qk_round_n, + params.l1_kv_pingpong_flag, params.hidden_size, is_rope); +} + +// 业务函数：QK MMAD 计算 +__aicore__ __attribute__((always_inline)) inline void ComputeQKMMad( + const QKParams ¶ms, uint32_t embed_split_idx, bool is_rope) +{ + PlatformComputeQKMMad(embed_split_idx, params.embed_split_size, + m, params.qk_n, params.qk_round_n_l1, + params.l1_kv_pingpong_flag, is_rope); +} + +// 业务函数：QK 结果拷贝到 GM +__aicore__ __attribute__((always_inline)) inline void CopyQKResultToGM( + const QKParams ¶ms, uint32_t embed_split_idx, + uint32_t n_idx, bool is_rope) +{ + PlatformCopyQKResultToGM(embed_split_idx, m, params.qk_n, + params.qk_round_n, params.l1_kv_pingpong_flag, + n_idx, is_rope); +} + +// 业务函数：初始化 PV 参数 +__aicore__ __attribute__((always_inline)) inline void InitPVParams( + MLAContext &ctx, uint32_t n_idx, PVParams ¶ms) +{ + params.qk_n_2 = ctx.qk_n_2; + params.qk_round_n_2 = ctx.qk_round_n_2; + params.qk_round_n_2_l1 = ctx.qk_round_n_2_l1; + params.k_round_n = ctx.k_round_n; + params.row_num = ctx.row_num; + params.hidden_size = ctx.hidden_size; + + if (n_idx == ctx.n_loop) { + params.qk_n_2 = (ctx.cur_kv_seqlen - (n_idx - 1) * ctx.pp_n_scalar); + params.qk_round_n_2 = RoundUp(params.qk_n_2); + params.qk_round_n_2_l1 = RoundUp(params.qk_n_2); + } + params.k_round_n = params.qk_round_n_2_l1; + params.l1_kv_pingpong_flag = (n_idx - 1) % 2; + params.l0_p_pingpong_flag = (n_idx - 1) % 2; + params.embed_split_size = 128; + embed_split_loop_v = 4; + params.round_embed_split_size = RoundUp(params.embed_split_size); +} + +// 业务函数：KV L1→L0B 转置加载 +__aicore__ __attribute__((always_inline)) inline void LoadKVTransposeToL0B( + const PVParams ¶ms, uint32_t embed_split_idx) +{ + bool is_last_split = (embed_split_idx == embed_split_loop_v - 1); + if (params.k_round_n <= params.round_embed_split_size) { + PlatformLoadKVTransposeSmallN( + params.l0b_pingpong_flag, params.l1_kv_pingpong_flag, + params.l1kv_offset, params.k_round_n, params.hidden_size, + params.embed_split_size, params.round_embed_split_size); + } else { + PlatformLoadKVTransposeLargeN( + params.l0b_pingpong_flag, params.l1_kv_pingpong_flag, + params.l1kv_offset, params.hidden_size, + params.qk_round_n_2, params.round_embed_split_size); + } + PlatformSetKVLoadComplete(is_last_split, params.l1_kv_pingpong_flag); +} + +// 业务函数：P 数据加载（GM→L1→L0A） +__aicore__ __attribute__((always_inline)) inline void LoadPDataToL0A( + const PVParams ¶ms, uint32_t n_idx) +{ + PlatformLoadPFromGMToL1( + n_idx, params.row_num, params.k_round_n, params.qk_round_n_2); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + PlatformLoadPToL0AInt8( + params.l0_p_pingpong_flag, params.row_num, + params.qk_round_n_2_l1); + } else { + PlatformLoadPToL0AGeneral( + params.l0_p_pingpong_flag, params.row_num, + params.qk_round_n_2); + } + PlatformSetPLoadComplete(); +} + +// 业务函数：PV MMAD 计算 +__aicore__ __attribute__((always_inline)) inline void ComputePVMmad( + const PVParams ¶ms, uint32_t embed_split_idx) +{ + bool is_last_split = (embed_split_idx == embed_split_loop_v - 1); + PlatformComputePVMmad( + params.l0b_pingpong_flag, params.l0c_pingpong_flag, + params.l0_p_pingpong_flag, params.embed_split_size, + params.qk_n_2, is_last_split); +} + +// 业务函数：PV 结果拷贝到 GM +__aicore__ __attribute__((always_inline)) inline void CopyPVResultToGM( + const PVParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) +{ + PlatformCopyPVResultToGM( + params.l0c_pingpong_flag, embed_split_idx, + n_idx, params.round_embed_split_size); +} + +// ==================== Rope 业务函数（INT8 专有）==================== + +// 业务函数：Rope 计算（Q_Rope × KV_Rope） +__aicore__ __attribute__((always_inline)) inline void ComputeQRope( + const QKParams ¶ms, uint32_t n_idx) +{ + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + // Rope 部分固定使用 embed_split_idx=4, embed_split_size=64 + uint32_t rope_embed_split_idx = 4; + uint32_t rope_embed_split_size = 64; + uint32_t rope_round_embed_split_size = 64; + + // 1. 加载 Q_Rope 到 L0A + PlatformLoadQRopeToL0A( + rope_embed_split_idx, params.q_load_coeff, + rope_round_embed_split_size); + + // 2. 加载 KV_Rope 到 L0B + PlatformLoadKVRopeToL0B( + rope_embed_split_idx, params.l1_kv_pingpong_flag, + rope_round_embed_split_size, params.qk_round_n); + + // 3. mmad 矩阵乘 + PlatformComputeQRopeMMad( + rope_embed_split_idx, params.l1_kv_pingpong_flag, + m, params.qk_n, rope_embed_split_size); + + // 4. 结果搬回 GM + PlatformCopyQRopeResultToGM( + params.l1_kv_pingpong_flag, n_idx, + m, params.qk_round_n); + } +} \ No newline at end of file From 1e9e1ec3960ef3f3bc23ed697184b790170ec4e5 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Tue, 11 Aug 2026 11:10:12 +0800 Subject: [PATCH 10/24] feat: you hua mla arch32 aiv --- .../op_kernel/multi_latent_attention.h | 714 +++---------- .../multi_latent_attention_aiv_arch32.h | 39 + .../op_kernel/multi_latent_attention_aiv_bs.h | 40 + .../op_kernel/multi_latent_attention_arch32.h | 195 ++++ .../op_kernel/multi_latent_attention_bs.h | 969 ++++++++++++++++++ 5 files changed, 1400 insertions(+), 557 deletions(-) create mode 100644 xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h create mode 100644 xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h index 0c9e62f..1d8b7fd 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h @@ -368,325 +368,16 @@ class MLAttentionDecoderAic { __aicore__ __attribute__((always_inline)) inline void Run() { - SET_FLAG(M, MTE1, EVENT_ID0); - SET_FLAG(M, MTE1, EVENT_ID1); - SET_FLAG(M, MTE1, EVENT_ID2); - SET_FLAG(M, MTE1, EVENT_ID3); - SET_FLAG(M, MTE1, EVENT_ID4); - SET_FLAG(M, MTE1, EVENT_ID5); - SET_FLAG(M, MTE1, EVENT_ID6); - SET_FLAG(M, MTE1, EVENT_ID7); - SET_FLAG(FIX, M, EVENT_ID0); - SET_FLAG(FIX, M, EVENT_ID1); - SET_FLAG(MTE1, MTE2, EVENT_ID0); - SET_FLAG(MTE1, MTE2, EVENT_ID1); - SET_FLAG(MTE1, MTE2, EVENT_ID2); - SET_FLAG(MTE1, MTE2, EVENT_ID3); - SET_FLAG(MTE1, MTE2, EVENT_ID4); - SET_FLAG(MTE1, MTE2, EVENT_ID5); - SET_FLAG(MTE1, MTE2, EVENT_ID6); - SET_FLAG(MTE1, MTE2, EVENT_ID7); - SET_FLAG(FIX, MTE1, EVENT_ID0); - SET_FLAG(FIX, MTE1, EVENT_ID1); - SET_FLAG(FIX, MTE1, EVENT_ID2); - SET_FLAG(FIX, MTE1, EVENT_ID3); - SET_FLAG(FIX, MTE1, EVENT_ID4); - SET_FLAG(FIX, MTE1, EVENT_ID5); - SET_FLAG(MTE2, FIX, EVENT_ID0); - - - uint64_t cur_batch = 0; - - uint32_t q_block_num_per_batch = (q_heads + cur_qn_blk_size - 1) / cur_qn_blk_size; - uint32_t process_num = q_block_num_per_batch * num_batches; - - for (uint32_t process = block_idx; process < process_num; process += (uint32_t)block_num) { // for task - cur_batch = process / q_block_num_per_batch; - if (cur_batch >= num_batches) break; - - uint32_t offset_tiling = tiling_head_size + tiling_para_size * cur_batch; - uint32_t start_core_idx = (cur_batch * q_block_num_per_batch) % block_num; - - uint32_t q_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); - if (kv_seqlen == 0) { - continue; - } - uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; - - uint32_t start_head = (process % q_block_num_per_batch) * cur_qn_blk_size; - uint32_t start_kv = 0; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen; - uint32_t cur_head_num = cur_qn_blk_size; - - InnerRunCubeMLA(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, - offset_tiling); - } - WAIT_FLAG(M, MTE1, EVENT_ID0); - WAIT_FLAG(M, MTE1, EVENT_ID1); - WAIT_FLAG(M, MTE1, EVENT_ID2); - WAIT_FLAG(M, MTE1, EVENT_ID3); - WAIT_FLAG(M, MTE1, EVENT_ID4); - WAIT_FLAG(M, MTE1, EVENT_ID5); - WAIT_FLAG(M, MTE1, EVENT_ID6); - WAIT_FLAG(M, MTE1, EVENT_ID7); - WAIT_FLAG(FIX, M, EVENT_ID0); - WAIT_FLAG(FIX, M, EVENT_ID1); - WAIT_FLAG(MTE1, MTE2, EVENT_ID0); - WAIT_FLAG(MTE1, MTE2, EVENT_ID1); - WAIT_FLAG(MTE1, MTE2, EVENT_ID2); - WAIT_FLAG(MTE1, MTE2, EVENT_ID3); - WAIT_FLAG(MTE1, MTE2, EVENT_ID4); - WAIT_FLAG(MTE1, MTE2, EVENT_ID5); - WAIT_FLAG(MTE1, MTE2, EVENT_ID6); - WAIT_FLAG(MTE1, MTE2, EVENT_ID7); - WAIT_FLAG(FIX, MTE1, EVENT_ID0); - WAIT_FLAG(FIX, MTE1, EVENT_ID1); - WAIT_FLAG(FIX, MTE1, EVENT_ID2); - WAIT_FLAG(FIX, MTE1, EVENT_ID3); - WAIT_FLAG(FIX, MTE1, EVENT_ID4); - WAIT_FLAG(FIX, MTE1, EVENT_ID5); - WAIT_FLAG(MTE2, FIX, EVENT_ID0); - PIPE_BARRIER(ALL); + PlatformInitPipeSync(); + ScheduleCubeTasks(); + PlatformWaitPipeSync(); } __aicore__ __attribute__((always_inline)) inline void RunTP1() { - SET_FLAG(M, MTE1, EVENT_ID0); - SET_FLAG(M, MTE1, EVENT_ID1); - SET_FLAG(M, MTE1, EVENT_ID2); - SET_FLAG(M, MTE1, EVENT_ID3); - SET_FLAG(M, MTE1, EVENT_ID4); - SET_FLAG(M, MTE1, EVENT_ID5); - SET_FLAG(M, MTE1, EVENT_ID6); - SET_FLAG(M, MTE1, EVENT_ID7); - SET_FLAG(FIX, M, EVENT_ID0); - SET_FLAG(FIX, M, EVENT_ID1); - SET_FLAG(MTE1, MTE2, EVENT_ID0); - SET_FLAG(MTE1, MTE2, EVENT_ID1); - SET_FLAG(MTE1, MTE2, EVENT_ID2); - SET_FLAG(MTE1, MTE2, EVENT_ID3); - SET_FLAG(MTE1, MTE2, EVENT_ID4); - SET_FLAG(MTE1, MTE2, EVENT_ID5); - SET_FLAG(MTE1, MTE2, EVENT_ID6); - SET_FLAG(MTE1, MTE2, EVENT_ID7); - SET_FLAG(FIX, MTE1, EVENT_ID0); - SET_FLAG(FIX, MTE1, EVENT_ID1); - SET_FLAG(FIX, MTE1, EVENT_ID2); - SET_FLAG(FIX, MTE1, EVENT_ID3); - SET_FLAG(FIX, MTE1, EVENT_ID4); - SET_FLAG(FIX, MTE1, EVENT_ID5); - SET_FLAG(MTE2, FIX, EVENT_ID0); - - uint32_t tail = totalTaskNum % block_num; - if constexpr (EnableOptimization) { - - } else{ - tail = 0; // control whether to run tail optimization - } - uint32_t totalTaskNumRound = totalTaskNum - tail; - - - for (uint32_t process = block_idx; process < totalTaskNumRound; process += (uint32_t)block_num) { // for task - uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; - uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - - uint32_t q_seqlen = 1; - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); - if (kv_seqlen == 0) { - continue; - } - uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; - - uint32_t start_head = 0; - uint32_t start_kv = 0; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen; - uint32_t cur_head_num = q_heads; - - InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); - } - - // suppose all seqs have same length - if (tail > 0){ - uint32_t sample_kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + tiling_head_size + 2)); - bool enableExtraOptimization = true; - if (block_num % 4 == 3) { - // cannot optimize this situation due to math problem - enableExtraOptimization = false; - } - if (sample_kv_seqlen <= 2048){ - // Too Short to benefit from optimization - enableExtraOptimization = false; - } - if (!enableExtraOptimization || tail <= block_num / 2) { - // collect all metadata - uint32_t cores_per_seq = 1; - if (0 < tail && tail <= block_num / 4) {// 6 tasks left, each works with 4 cores - cores_per_seq = 4; - if (tail == 1){ - cores_per_seq = block_num; - } - else if (tail == 2){ - cores_per_seq = block_num / 2; - } - else if(tail == 3){ - cores_per_seq = block_num / 3; - } - else if(tail == 4){ - cores_per_seq = block_num / 4; - } - } - else if(block_num / 4 < tail && tail <= block_num / 3) { // 8 tasks left, each works with 3 cores - cores_per_seq = 3; - - } - else if(block_num / 3 < tail && tail <= block_num / 2) { // 12 tasks left, each works with 2 cores - cores_per_seq = 2; - } - else { - // no extra optimization for tail > 12 - cores_per_seq = 1; - } - - if(!enableExtraOptimization){ - cores_per_seq = 1; - } - - uint32_t process = totalTaskNumRound + block_idx / cores_per_seq; - if (process < totalTaskNum) { - uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; - uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - - uint32_t q_seqlen = 1; - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); - uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; - uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; - uint32_t actual_work_cores = kv_seqlen / kv_seqlen_align + (kv_seqlen % kv_seqlen_align != 0); - // cores_per_seq = actual_work_cores; - uint32_t kv_seqlen_process = 0; - if (block_idx < block_idx / cores_per_seq * cores_per_seq + actual_work_cores){ - kv_seqlen_process = (block_idx % cores_per_seq == actual_work_cores - 1) ? - (kv_seqlen - kv_seqlen_align * (actual_work_cores - 1)) : kv_seqlen_align; - } - - if (kv_seqlen > 0 && kv_seqlen_process > 0) { - uint32_t start_head = 0; - uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen_process; - uint32_t cur_head_num = q_heads; - - // no need to modify anything in cube kernel, just call the same kernel - InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); - } - } - } - else if (tail > 3 * block_num / 4){ - // no benefit for optimizing this situation - uint32_t process = totalTaskNumRound + block_idx; - if (process < totalTaskNum) { - uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; - uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - - uint32_t q_seqlen = 1; - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); - if (kv_seqlen > 0) { - uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; - - uint32_t start_head = 0; - uint32_t start_kv = 0; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen; - uint32_t cur_head_num = q_heads; - - InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); - } - } - } - else { - // 18 >= tail >= 12 - // first 12 tasks, two cores per task - { - uint32_t cores_per_seq = 2; - uint32_t process = totalTaskNumRound + block_idx / cores_per_seq; - uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; - uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - - uint32_t q_seqlen = 1; - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); - uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; - uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; - uint32_t kv_seqlen_process = (block_idx % cores_per_seq == cores_per_seq - 1) ? - (kv_seqlen - kv_seqlen_align * (cores_per_seq - 1)) : kv_seqlen_align; - - if (kv_seqlen > 0 && kv_seqlen_process > 0) { - uint32_t start_head = 0; - uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen_process; - uint32_t cur_head_num = q_heads; - - // no need to modify anything in cube kernel, just call the same kernel - InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); - } - } - { - uint32_t cores_per_seq = 4; - uint32_t process = totalTaskNumRound + block_num / 2 + block_idx / cores_per_seq; - if (process < totalTaskNum) { - uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; - uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - - uint32_t q_seqlen = 1; - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); - uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; - uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; - uint32_t kv_seqlen_process = (block_idx % cores_per_seq == cores_per_seq - 1) ? - (kv_seqlen - kv_seqlen_align * (cores_per_seq - 1)) : kv_seqlen_align; - - if (kv_seqlen > 0 && kv_seqlen_process > 0) { - uint32_t start_head = 0; - uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen_process; - uint32_t cur_head_num = q_heads; - - // no need to modify anything in cube kernel, just call the same kernel - InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); - } - } - } - } - } - - WAIT_FLAG(M, MTE1, EVENT_ID0); - WAIT_FLAG(M, MTE1, EVENT_ID1); - WAIT_FLAG(M, MTE1, EVENT_ID2); - WAIT_FLAG(M, MTE1, EVENT_ID3); - WAIT_FLAG(M, MTE1, EVENT_ID4); - WAIT_FLAG(M, MTE1, EVENT_ID5); - WAIT_FLAG(M, MTE1, EVENT_ID6); - WAIT_FLAG(M, MTE1, EVENT_ID7); - WAIT_FLAG(FIX, M, EVENT_ID0); - WAIT_FLAG(FIX, M, EVENT_ID1); - WAIT_FLAG(MTE1, MTE2, EVENT_ID0); - WAIT_FLAG(MTE1, MTE2, EVENT_ID1); - WAIT_FLAG(MTE1, MTE2, EVENT_ID2); - WAIT_FLAG(MTE1, MTE2, EVENT_ID3); - WAIT_FLAG(MTE1, MTE2, EVENT_ID4); - WAIT_FLAG(MTE1, MTE2, EVENT_ID5); - WAIT_FLAG(MTE1, MTE2, EVENT_ID6); - WAIT_FLAG(MTE1, MTE2, EVENT_ID7); - WAIT_FLAG(FIX, MTE1, EVENT_ID0); - WAIT_FLAG(FIX, MTE1, EVENT_ID1); - WAIT_FLAG(FIX, MTE1, EVENT_ID2); - WAIT_FLAG(FIX, MTE1, EVENT_ID3); - WAIT_FLAG(FIX, MTE1, EVENT_ID4); - WAIT_FLAG(FIX, MTE1, EVENT_ID5); - WAIT_FLAG(MTE2, FIX, EVENT_ID0); - PIPE_BARRIER(ALL); + PlatformInitPipeSync(); + ScheduleCubeTasksTP1(); + PlatformWaitPipeSync(); } private: @@ -902,11 +593,137 @@ class MLAttentionDecoderAic { uint64_t l1kv_offset; }; + // ==================== InnerRunCubeMLATP1 三层架构拆分 ==================== + // === 第一层：TP1Context 结构体 === + struct TP1Context { + // input params + uint32_t cur_batch; + uint32_t start_head; + uint32_t cur_head_num; + uint32_t start_kv; + uint32_t cur_q_seqlen; + uint32_t cur_kv_seqlen; + uint32_t offset_tiling; + + // Q address + uint64_t q_offset; + uint64_t q_rope_offset; + + // loop & size + uint32_t pp_n_scalar; + uint32_t sub_n_loop; + uint32_t n_loop; + + // QK dims + uint32_t qk_n; + uint32_t qk_round_n; + uint32_t qk_n_2; + uint32_t qk_round_n_2; + + // row info + uint32_t row_num; + // m is class member + + // TP1 专有字段 + uint32_t sv_n; // PV阶段有效列数（动态变化） + uint32_t s_block_stack; // 流水线深度=4 + }; + + // === 第一层：InitTP1Context 编排层 === + __aicore__ __attribute__((always_inline)) inline void InitTP1Context( + TP1Context &ctx, uint32_t cur_batch, uint32_t start_head, uint32_t cur_head_num, + uint32_t start_kv, uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t offset_tiling) + { + ctx.cur_batch = cur_batch; + ctx.start_head = start_head; + ctx.cur_head_num = cur_head_num; + ctx.start_kv = start_kv; + ctx.cur_q_seqlen = cur_q_seqlen; + ctx.cur_kv_seqlen = cur_kv_seqlen; + ctx.offset_tiling = offset_tiling; + + uint32_t prev_task = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); + uint64_t addr_q_scalar = (uint64_t)prev_task * q_heads; + ctx.q_offset = addr_q_scalar * 512 + start_head * 512; + ctx.q_rope_offset = addr_q_scalar * 64 + start_head * 64; + + ctx.pp_n_scalar = block_size; + ctx.sub_n_loop = ctx.pp_n_scalar / block_size; + ctx.n_loop = (cur_kv_seqlen + ctx.pp_n_scalar - 1) / ctx.pp_n_scalar; + + ctx.qk_n = ctx.pp_n_scalar; + ctx.qk_round_n = RoundUp(ctx.qk_n); + ctx.qk_n_2 = ctx.pp_n_scalar; + ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); + + ctx.row_num = cur_head_num * cur_q_seqlen; + ctx.sv_n = ctx.n_loop == 1 ? cur_kv_seqlen : ctx.pp_n_scalar; + m = RoundUp<16>(ctx.row_num); + + ctx.s_block_stack = 4; + } + + // === 第一层：LoadTP1QData 编排层 === + __aicore__ __attribute__((always_inline)) inline void LoadTP1QData(TP1Context &ctx) + { + LoadQMainFromGMToL1(l1q_buf_addr_tensor, q_gm_tensor, ctx.q_offset, ctx.cur_q_seqlen, ctx.cur_head_num); + LoadQRopeFromGMToL1(l1q_buf_addr_tensor, l1q_rope_buf_addr_tensor, q_rope_gm_tensor, ctx.q_rope_offset, ctx.cur_q_seqlen, ctx.cur_head_num); + PlatformSetQLoadComplete(); + } + + // [PV 第三层平台函数 + 第二层业务函数] 已迁移至 arch32.h / bs.h - // === 以下 include 第三层和第二层子文件（放在 QKParams/PVParams 之后，确保类型可见）=== + // === 以下 include 第三层和第二层子文件（放在 QKParams/PVParams/TP1Context 之后，确保类型可见）=== #include "multi_latent_attention_arch32.h" #include "multi_latent_attention_bs.h" + // === 第一层：ComputeTP1QK 编排层（CUBE1 stage1）=== + // 外层 split_idx 循环 + 内层 embed_split_idx 循环（5次: 0-4） + // idx 0,2: Q→L0A + KV Main GM→L1 + // idx 1,3: Q→L0A + KV L1→L0B + mmad + // idx 4: Q→L0A + KV Rope GM→L1 + L1→L0B + mmad + S→GM + // 循环后: FftsCrossCoreSync + __aicore__ __attribute__((always_inline)) inline void ComputeTP1QK( + TP1Context &ctx, uint32_t n_idx) + { + for (uint32_t split_idx = 0; split_idx < ctx.s_block_stack && n_idx + split_idx < ctx.n_loop; split_idx++) { + TP1QKParams params; + InitTP1QKParams(ctx, split_idx, n_idx, params); + + for (uint32_t embed_split_idx = 0; embed_split_idx < 5; ++embed_split_idx) { + if (embed_split_idx == 4) { + params.embed_split_size = 64; + params.round_embed_split_size = 64; + } + + // Q L1→L0A + LoadTP1QDataToL0A(params, embed_split_idx); + + // KV Main GM→L1 (idx 0,2) + if (embed_split_idx == 0 || embed_split_idx == 2) { + LoadTP1KVMainFromGM(params, embed_split_idx); + } + + // KV Rope GM→L1 (idx 4) + if (embed_split_idx == 4) { + LoadTP1KVRopeFromGM(params, embed_split_idx); + } + + // KV L1→L0B (all idx) + LoadTP1KVDataToL0B(params, embed_split_idx); + + // QK mmad + ComputeTP1QKMMad(params, embed_split_idx); + + // S→GM (idx 4) + if (embed_split_idx == 4) { + CopyTP1QKResultToGM(params, embed_split_idx, n_idx); + } + } + } + FftsCrossCoreSync(QK_READY_DECODER); + } + // === 第一层：ComputePV 编排层 === __aicore__ __attribute__((always_inline)) inline void ComputePV( MLAContext &ctx, uint32_t n_idx) @@ -967,204 +784,27 @@ class MLAttentionDecoderAic { __aicore__ __attribute__((always_inline)) inline void InnerRunCubeMLATP1(uint32_t cur_batch, uint32_t start_head, uint32_t cur_head_num, uint32_t start_kv, uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t offset_tiling) { - uint32_t prev_task = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); - uint64_t addr_q_scalar = (uint64_t)prev_task * q_heads; - uint64_t q_offset = addr_q_scalar * 512 + start_head * 512; - uint64_t q_rope_offset = addr_q_scalar * 64 + start_head * 64; - - uint32_t pp_n_scalar = block_size; - uint32_t sub_n_loop = pp_n_scalar / block_size; - - uint32_t n_loop = (cur_kv_seqlen + pp_n_scalar - 1) / pp_n_scalar; - - uint32_t qk_n = pp_n_scalar; - uint32_t qk_round_n = RoundUp(qk_n); - uint32_t qk_n_2 = pp_n_scalar; - uint32_t qk_round_n_2 = RoundUp(qk_n_2); - - uint32_t row_num = cur_head_num * cur_q_seqlen; - - uint32_t sv_n = n_loop == 1 ? cur_kv_seqlen : pp_n_scalar; - m = RoundUp<16>(row_num); - - // copy Q - LoadQMainFromGMToL1(l1q_buf_addr_tensor, q_gm_tensor, q_offset, cur_q_seqlen, cur_head_num); - LoadQRopeFromGMToL1(l1q_buf_addr_tensor, l1q_rope_buf_addr_tensor, q_rope_gm_tensor, q_rope_offset, cur_q_seqlen, cur_head_num); + TP1Context ctx; + InitTP1Context(ctx, cur_batch, start_head, cur_head_num, + start_kv, cur_q_seqlen, cur_kv_seqlen, offset_tiling); + LoadTP1QData(ctx); - SET_FLAG(MTE2, MTE1, EVENT_ID0); - WAIT_FLAG(MTE2, MTE1, EVENT_ID0); - uint32_t s_block_stack = 4; + uint32_t pp_n_scalar = ctx.pp_n_scalar; + uint32_t n_loop = ctx.n_loop; + uint32_t s_block_stack = ctx.s_block_stack; for (uint32_t n_idx = 0; n_idx < n_loop + s_block_stack; n_idx+=s_block_stack) { if (n_idx < n_loop) { - uint32_t sv_n_triu = n_loop * pp_n_scalar; - if (n_idx + s_block_stack > n_loop - 1) { - sv_n = cur_kv_seqlen - n_idx * pp_n_scalar; // delete - } else { - sv_n = pp_n_scalar * s_block_stack; - } - uint32_t sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; - for (uint32_t split_idx = 0; split_idx < s_block_stack && n_idx + split_idx < n_loop; split_idx++) { - uint32_t now_idx = n_idx + split_idx; - uint32_t l1_kv_pingpong_flag = now_idx % 2; - if (now_idx == (n_loop - 1)) { - qk_n = (cur_kv_seqlen - now_idx * pp_n_scalar); - qk_round_n = RoundUp(qk_n); - } else { - qk_n = pp_n_scalar; - qk_round_n = RoundUp(qk_n); - } - bool last_split = split_idx == s_block_stack - 1 || now_idx == n_loop - 1; - uint32_t embed_split_size = 128; - uint32_t round_embed_split_size = RoundUp(embed_split_size); - - /* ************ CUBE1 stage1 ************* */ - - uint32_t block_table_id = (uint32_t)(*(block_tables_gm + - cur_batch * max_num_blocks_per_query + start_kv / block_size + now_idx)); - int64_t kv_offset = (int64_t)block_table_id * block_size * stride_kv; - int64_t kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; - - - uint32_t q_load_coeff = 1; - q_load_coeff = m; - int64_t now_l1_offset = 0; - for (uint32_t embed_split_idx = 0; embed_split_idx < 5; ++embed_split_idx) { - if (embed_split_idx == 4) { - embed_split_size = 64; - round_embed_split_size = 64; - } - WAIT_FLAG(M, MTE1, embed_split_idx % 2); - - for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { - l1_to_l0_a( - l0a_buf_tensor[embed_split_idx % 2 * 16384 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], - l1q_buf_addr_tensor[embed_split_idx * m * 128 + loa_load_idx * T_CUBE_MATRIX_SIZE], - 0, - round_embed_split_size / T_BLOCK_SIZE, // repeat - 0, - q_load_coeff / BLOCK_SIZE, // srcStride - 0, - 0 // dstStride - ); - } - - SET_FLAG(MTE1, M, embed_split_idx % 2); - if (embed_split_idx == 0 || embed_split_idx == 2) { - WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); // 等待V全部搬入L0B - now_l1_offset = l1_kv_pingpong_flag * 128 * 256; - if constexpr(KInputType == InputFormat::ND_FORMAT) { - gm_to_l1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 256], - k_gm_tensor[kv_offset + embed_split_idx * 128], - qk_n, // nValue - qk_round_n, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 256, // dValue - 0, // dstNzMatrixStride, unused - stride_kv // srcDValue - ); - } else { - gm_to_l1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 256], - k_gm_tensor[kv_offset + block_size * 128 * embed_split_idx], - qk_round_n, - block_size, - qk_round_n, - 256, - 256, 256); - } - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - } else if (embed_split_idx == 4) { - WAIT_FLAG(MTE1, MTE2, 2 + l1_kv_pingpong_flag); - now_l1_offset = l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; - if constexpr(KInputType == InputFormat::ND_FORMAT) { - gm_to_l1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128], - k_rope_gm_tensor[kv_offset_rope], - qk_n, // nValue - qk_round_n, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 64, // dValue - 0, // dstNzMatrixStride, unused - stride_kv_rope // srcDValue - ); - } else { - gm_to_l1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128], - k_rope_gm_tensor[kv_offset_rope], - qk_round_n, - block_size, - qk_round_n, - 64, - 64, - 64); - } - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - } - WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); - l1_to_l0_b( - l0b_buf_tensor[embed_split_idx % 2 * 16384], - l1kv_buf_addr_tensor[now_l1_offset + embed_split_idx % 2 * qk_round_n * 128], - 0, - round_embed_split_size * qk_round_n / T_CUBE_MATRIX_SIZE, // repeat - 0, - 1, // srcStride - 0, - 0 // dstStride - ); - if (embed_split_idx == 1 || embed_split_idx == 3) { - SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); - } - if (embed_split_idx == 4) { - SET_FLAG(MTE1, MTE2, 2 + l1_kv_pingpong_flag); - } - SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); - WAIT_FLAG(MTE1, M, embed_split_idx % 2); - WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); - if (embed_split_idx == 0) { - WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); - } - mmad( - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - l0a_buf_tensor[embed_split_idx % 2 * 16384], - l0b_buf_tensor[embed_split_idx % 2 * 16384], - m, // m - qk_n, // n - embed_split_size, // k - embed_split_idx == 0 // cmatrixInitVal - ); - PIPE_BARRIER(M); - SET_FLAG(M, MTE1, embed_split_idx % 2); - SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); - - // copy S to gm - if (embed_split_idx == 4) { - SET_FLAG(M, FIX, l1_kv_pingpong_flag); - WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); - - l0c_to_gm( - s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER * 4 + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar], - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - m, // MSize - qk_round_n, // NSize - RoundUp<16>(m), // srcStride - sv_round_n // dstStride_dst_D - ); - SET_FLAG(FIX, M, l1_kv_pingpong_flag); - } - } - } - FftsCrossCoreSync(QK_READY_DECODER); + ComputeTP1QK(ctx, n_idx); } /* ************ CUBE2 stage1 ************* */ if (n_idx >= s_block_stack) { if (n_idx + s_block_stack > n_loop + s_block_stack - 1) { - sv_n = cur_kv_seqlen - (n_idx - s_block_stack) * pp_n_scalar; // delete + ctx.sv_n = cur_kv_seqlen - (n_idx - s_block_stack) * pp_n_scalar; // delete } else { - sv_n = pp_n_scalar * s_block_stack; + ctx.sv_n = pp_n_scalar * s_block_stack; } + uint32_t sv_n = ctx.sv_n; + uint32_t row_num = ctx.row_num; uint32_t sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; uint32_t embed_split_size = 128; embed_split_loop_v = 4; @@ -1175,11 +815,13 @@ class MLAttentionDecoderAic { for (uint32_t split_idx = 0; split_idx < s_block_stack && n_idx + split_idx < n_loop + s_block_stack; split_idx++) { uint32_t now_idx = n_idx + split_idx; if (now_idx == (n_loop + s_block_stack - 1)) { - qk_n_2 = (cur_kv_seqlen - (now_idx - s_block_stack) * pp_n_scalar); + ctx.qk_n_2 = (cur_kv_seqlen - (now_idx - s_block_stack) * pp_n_scalar); } else { - qk_n_2 = pp_n_scalar; + ctx.qk_n_2 = pp_n_scalar; } - qk_round_n_2 = RoundUp(qk_n_2); + ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); + uint32_t qk_n_2 = ctx.qk_n_2; + uint32_t qk_round_n_2 = ctx.qk_round_n_2; uint32_t l1_kv_pingpong_flag = now_idx % 2; uint32_t l0_p_pingpong_flag = now_idx % 2; uint32_t l0b_pingpong_flag = now_idx % 2; @@ -1485,55 +1127,9 @@ class MLADecoderAiv{ __aicore__ __attribute__((always_inline)) inline void Run() { - SET_FLAG(MTE3, V, EVENT_ID0); - SET_FLAG(MTE3, MTE2, EVENT_ID0); - SET_FLAG(MTE3, MTE2, EVENT_ID2); - SET_FLAG(MTE3, MTE2, EVENT_ID3); - SET_FLAG(MTE3, MTE2, EVENT_ID4); - SET_FLAG(V, MTE2, EVENT_ID4); - SET_FLAG(V, MTE2, EVENT_ID0); - SET_FLAG(MTE3, V, EVENT_ID2); - SET_FLAG(V, MTE2, EVENT_ID2); - - - uint64_t cur_batch = 0; - - uint32_t q_block_num_per_batch = (q_heads + cur_qn_blk_size - 1) / cur_qn_blk_size; - uint32_t process_num = q_block_num_per_batch * num_batches; - - for (uint32_t process = block_idx; process < process_num; process += (uint32_t)block_num) { // for task - cur_batch = process / q_block_num_per_batch; - if (cur_batch >= num_batches) break; - - uint32_t offset_tiling = tiling_head_size + tiling_para_size * cur_batch; - uint32_t start_core_idx = (cur_batch * q_block_num_per_batch) % block_num; - - uint32_t q_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); - if (kv_seqlen == 0) { - continue; - } - uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; - - uint32_t start_head = (process % q_block_num_per_batch) * cur_qn_blk_size; - uint32_t start_kv = 0; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen; - uint32_t cur_head_num = cur_qn_blk_size; - uint32_t cur_nIndx = 0; - InnerRunVectorChange(cur_batch, start_head, cur_nIndx, cur_q_seq_len, cur_kv_seqlen, cur_head_num, - offset_tiling, 512, embed_split_loop_v_former); - } - - WAIT_FLAG(MTE3, V, EVENT_ID0); - WAIT_FLAG(MTE3, MTE2, EVENT_ID0); - WAIT_FLAG(MTE3, MTE2, EVENT_ID2); - WAIT_FLAG(MTE3, MTE2, EVENT_ID3); - WAIT_FLAG(MTE3, MTE2, EVENT_ID4); - WAIT_FLAG(V, MTE2, EVENT_ID0); - WAIT_FLAG(V, MTE2, EVENT_ID4); - WAIT_FLAG(MTE3, V, EVENT_ID2); - WAIT_FLAG(V, MTE2, EVENT_ID2); + PlatformInitVectorPipeSync(); + ScheduleVectorTasks(); + PlatformWaitVectorPipeSync(); } __aicore__ __attribute__((always_inline)) inline void RunTP1() @@ -1808,6 +1404,10 @@ class MLADecoderAiv{ } private: + // [AIV 第三层平台函数 + 第二层业务函数] 已迁移至 aiv_arch32.h / aiv_bs.h + #include "multi_latent_attention_aiv_arch32.h" + #include "multi_latent_attention_aiv_bs.h" + __aicore__ __attribute__((always_inline)) inline void ReduceMaxRepeatM( const AscendC::LocalTensor& dst, diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h new file mode 100644 index 0000000..91c6ca7 --- /dev/null +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h @@ -0,0 +1,39 @@ +// Copyright 2025 The xLLM Authors. 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: +// +// https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE +// +// ==================== AIV 平台处理层（第三层）==================== +// 本文件仅被 MLADecoderAiv 类 include，包含 Vector 侧平台同步函数。 +// AIC 侧平台函数在 multi_latent_attention_arch32.h 中。 + +// 平台函数：Vector 非TP1 管道同步初始化（Run 方法头部 9 条 SET_FLAG） +__aicore__ __attribute__((always_inline)) inline void PlatformInitVectorPipeSync() +{ + SET_FLAG(MTE3, V, EVENT_ID0); + SET_FLAG(MTE3, MTE2, EVENT_ID0); + SET_FLAG(MTE3, MTE2, EVENT_ID2); + SET_FLAG(MTE3, MTE2, EVENT_ID3); + SET_FLAG(MTE3, MTE2, EVENT_ID4); + SET_FLAG(V, MTE2, EVENT_ID4); + SET_FLAG(V, MTE2, EVENT_ID0); + SET_FLAG(MTE3, V, EVENT_ID2); + SET_FLAG(V, MTE2, EVENT_ID2); +} + +// 平台函数：等待 Vector 非TP1 管道同步完成（Run 方法尾部 9 条 WAIT_FLAG） +__aicore__ __attribute__((always_inline)) inline void PlatformWaitVectorPipeSync() +{ + WAIT_FLAG(MTE3, V, EVENT_ID0); + WAIT_FLAG(MTE3, MTE2, EVENT_ID0); + WAIT_FLAG(MTE3, MTE2, EVENT_ID2); + WAIT_FLAG(MTE3, MTE2, EVENT_ID3); + WAIT_FLAG(MTE3, MTE2, EVENT_ID4); + WAIT_FLAG(V, MTE2, EVENT_ID0); + WAIT_FLAG(V, MTE2, EVENT_ID4); + WAIT_FLAG(MTE3, V, EVENT_ID2); + WAIT_FLAG(V, MTE2, EVENT_ID2); +} \ No newline at end of file diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h new file mode 100644 index 0000000..215afc0 --- /dev/null +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h @@ -0,0 +1,40 @@ +// Copyright 2025 The xLLM Authors. 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: +// +// https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE +// +// ==================== AIV 业务逻辑层（第二层）==================== +// 本文件仅被 MLADecoderAiv 类 include，包含 Vector 侧业务调度函数。 +// AIC 侧业务函数在 multi_latent_attention_bs.h 中。 + +// 业务函数：调度 Vector 非 TP1 任务（Run 方法的中段业务逻辑） +// 包含 batch/head 循环调度，读取 tiling 参数后调用 InnerRunVectorChange +__aicore__ __attribute__((always_inline)) inline void ScheduleVectorTasks() +{ + uint64_t cur_batch = 0; + uint32_t q_block_num_per_batch = (q_heads + cur_qn_blk_size - 1) / cur_qn_blk_size; + uint32_t process_num = q_block_num_per_batch * num_batches; + for (uint32_t process = block_idx; process < process_num; process += (uint32_t)block_num) { + cur_batch = process / q_block_num_per_batch; + if (cur_batch >= num_batches) break; + uint32_t offset_tiling = tiling_head_size + tiling_para_size * cur_batch; + uint32_t start_core_idx = (cur_batch * q_block_num_per_batch) % block_num; + uint32_t q_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); + if (kv_seqlen == 0) { + continue; + } + uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; + uint32_t start_head = (process % q_block_num_per_batch) * cur_qn_blk_size; + uint32_t start_kv = 0; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen; + uint32_t cur_head_num = cur_qn_blk_size; + uint32_t cur_nIndx = 0; + InnerRunVectorChange(cur_batch, start_head, cur_nIndx, cur_q_seq_len, cur_kv_seqlen, cur_head_num, + offset_tiling, 512, embed_split_loop_v_former); + } +} \ No newline at end of file diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_arch32.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_arch32.h index a669a5e..19ebcc0 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_arch32.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_arch32.h @@ -442,4 +442,199 @@ __aicore__ __attribute__((always_inline)) inline void PlatformCopyQRopeResultToG ); } SET_FLAG(FIX, M, l1_kv_pingpong_flag); +} + +// 平台函数：初始化管道同步事件（Run 方法头部批量 SET_FLAG） +__aicore__ __attribute__((always_inline)) inline void PlatformInitPipeSync() +{ + SET_FLAG(M, MTE1, EVENT_ID0); + SET_FLAG(M, MTE1, EVENT_ID1); + SET_FLAG(M, MTE1, EVENT_ID2); + SET_FLAG(M, MTE1, EVENT_ID3); + SET_FLAG(M, MTE1, EVENT_ID4); + SET_FLAG(M, MTE1, EVENT_ID5); + SET_FLAG(M, MTE1, EVENT_ID6); + SET_FLAG(M, MTE1, EVENT_ID7); + SET_FLAG(FIX, M, EVENT_ID0); + SET_FLAG(FIX, M, EVENT_ID1); + SET_FLAG(MTE1, MTE2, EVENT_ID0); + SET_FLAG(MTE1, MTE2, EVENT_ID1); + SET_FLAG(MTE1, MTE2, EVENT_ID2); + SET_FLAG(MTE1, MTE2, EVENT_ID3); + SET_FLAG(MTE1, MTE2, EVENT_ID4); + SET_FLAG(MTE1, MTE2, EVENT_ID5); + SET_FLAG(MTE1, MTE2, EVENT_ID6); + SET_FLAG(MTE1, MTE2, EVENT_ID7); + SET_FLAG(FIX, MTE1, EVENT_ID0); + SET_FLAG(FIX, MTE1, EVENT_ID1); + SET_FLAG(FIX, MTE1, EVENT_ID2); + SET_FLAG(FIX, MTE1, EVENT_ID3); + SET_FLAG(FIX, MTE1, EVENT_ID4); + SET_FLAG(FIX, MTE1, EVENT_ID5); + SET_FLAG(MTE2, FIX, EVENT_ID0); +} + +// 平台函数：等待管道同步完成（Run 方法尾部批量 WAIT_FLAG + BARRIER） +__aicore__ __attribute__((always_inline)) inline void PlatformWaitPipeSync() +{ + WAIT_FLAG(M, MTE1, EVENT_ID0); + WAIT_FLAG(M, MTE1, EVENT_ID1); + WAIT_FLAG(M, MTE1, EVENT_ID2); + WAIT_FLAG(M, MTE1, EVENT_ID3); + WAIT_FLAG(M, MTE1, EVENT_ID4); + WAIT_FLAG(M, MTE1, EVENT_ID5); + WAIT_FLAG(M, MTE1, EVENT_ID6); + WAIT_FLAG(M, MTE1, EVENT_ID7); + WAIT_FLAG(FIX, M, EVENT_ID0); + WAIT_FLAG(FIX, M, EVENT_ID1); + WAIT_FLAG(MTE1, MTE2, EVENT_ID0); + WAIT_FLAG(MTE1, MTE2, EVENT_ID1); + WAIT_FLAG(MTE1, MTE2, EVENT_ID2); + WAIT_FLAG(MTE1, MTE2, EVENT_ID3); + WAIT_FLAG(MTE1, MTE2, EVENT_ID4); + WAIT_FLAG(MTE1, MTE2, EVENT_ID5); + WAIT_FLAG(MTE1, MTE2, EVENT_ID6); + WAIT_FLAG(MTE1, MTE2, EVENT_ID7); + WAIT_FLAG(FIX, MTE1, EVENT_ID0); + WAIT_FLAG(FIX, MTE1, EVENT_ID1); + WAIT_FLAG(FIX, MTE1, EVENT_ID2); + WAIT_FLAG(FIX, MTE1, EVENT_ID3); + WAIT_FLAG(FIX, MTE1, EVENT_ID4); + WAIT_FLAG(FIX, MTE1, EVENT_ID5); + WAIT_FLAG(MTE2, FIX, EVENT_ID0); + PIPE_BARRIER(ALL); +} + +// ==================== TP1 QK 平台函数 ==================== + +// 平台函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2 分支） +// TP1 专有：L1 目标偏移 [l1_kv_pingpong_flag * 128 * 256]，dValue=256，使用 block_table_id+embed_split_idx 寻址 +__aicore__ __attribute__((always_inline)) inline void PlatformLoadTP1KVMainToL1( + uint32_t embed_split_idx, uint32_t qk_n, uint32_t qk_round_n, + uint32_t l1_kv_pingpong_flag, int64_t kv_offset, uint32_t embed_split_idx_128) +{ + WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); + if constexpr (KInputType == InputFormat::ND_FORMAT) { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 256], + k_gm_tensor[kv_offset + embed_split_idx_128], + qk_n, + qk_round_n, + 0, + 256, + 0, + stride_kv); + } else { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 256], + k_gm_tensor[kv_offset + block_size * 128 * embed_split_idx], + qk_round_n, + block_size, + qk_round_n, + 256, + 256, 256); + } + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); +} + +// 平台函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4 分支） +// TP1 专有：L1 目标偏移 [l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128]，dValue=64 +__aicore__ __attribute__((always_inline)) inline void PlatformLoadTP1KVRopeToL1( + uint32_t qk_n, uint32_t qk_round_n, + uint32_t l1_kv_pingpong_flag, int64_t kv_offset_rope) +{ + WAIT_FLAG(MTE1, MTE2, 2 + l1_kv_pingpong_flag); + if constexpr (KInputType == InputFormat::ND_FORMAT) { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128], + k_rope_gm_tensor[kv_offset_rope], + qk_n, + qk_round_n, + 0, + 64, + 0, + stride_kv_rope); + } else { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128], + k_rope_gm_tensor[kv_offset_rope], + qk_round_n, + block_size, + qk_round_n, + 64, + 64, + 64); + } + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); +} + +// 平台函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） +// TP1 专有：L1 源偏移 [now_l1_offset + embed_split_idx % 2 * qk_round_n * 128] +__aicore__ __attribute__((always_inline)) inline void PlatformLoadTP1KVToL0B( + uint32_t embed_split_idx, uint32_t round_embed_split_size, + uint32_t qk_round_n, int64_t now_l1_offset, uint32_t l1_kv_pingpong_flag) +{ + WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); + l1_to_l0_b( + l0b_buf_tensor[embed_split_idx % 2 * 16384], + l1kv_buf_addr_tensor[now_l1_offset + embed_split_idx % 2 * qk_round_n * 128], + 0, + round_embed_split_size * qk_round_n / T_CUBE_MATRIX_SIZE, + 0, + 1, + 0, + 0); + // 释放 KV L1 缓冲区 + if (embed_split_idx == 1 || embed_split_idx == 3) { + SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); + } + if (embed_split_idx == 4) { + SET_FLAG(MTE1, MTE2, 2 + l1_kv_pingpong_flag); + } + SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); +} + +// 平台函数：TP1 QK MMAD 计算（含同步） +// TP1 与非 TP1 的 cmatrixInitVal 均为 (embed_split_idx == 0)，可复用逻辑 +__aicore__ __attribute__((always_inline)) inline void PlatformComputeTP1QKMMad( + uint32_t embed_split_idx, uint32_t embed_split_size, + uint32_t m_value, uint32_t qk_n, uint32_t l1_kv_pingpong_flag) +{ + WAIT_FLAG(MTE1, M, embed_split_idx % 2); + WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); + if (embed_split_idx == 0) { + WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); + } + mmad( + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + l0a_buf_tensor[embed_split_idx % 2 * 16384], + l0b_buf_tensor[embed_split_idx % 2 * 16384], + m_value, + qk_n, + embed_split_size, + embed_split_idx == 0); + PIPE_BARRIER(M); + SET_FLAG(M, MTE1, embed_split_idx % 2); + SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); +} + +// 平台函数：TP1 QK 结果从 L0C 拷贝到 GM（embed_split_idx == 4 时） +// TP1 专有：GM 目标偏移 [block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar] +// dstStride = sv_round_n +__aicore__ __attribute__((always_inline)) inline void PlatformCopyTP1QKResultToGM( + uint32_t m_value, uint32_t qk_round_n, uint32_t l1_kv_pingpong_flag, + uint64_t gm_dst_offset, uint32_t sv_round_n) +{ + SET_FLAG(M, FIX, l1_kv_pingpong_flag); + WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); + l0c_to_gm( + s_gm_tensor[gm_dst_offset], + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + m_value, + qk_round_n, + RoundUp<16>(m_value), + sv_round_n); + SET_FLAG(FIX, M, l1_kv_pingpong_flag); } \ No newline at end of file diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h index 8395eba..8b91ac4 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h @@ -82,6 +82,131 @@ __aicore__ __attribute__((always_inline)) inline void LoadQMainFromGMToL1( } } +// ==================== TP1 QK 业务函数 ==================== + +// TP1 QK 参数结构体 +struct TP1QKParams { + uint32_t qk_n; + uint32_t qk_round_n; + uint32_t l1_kv_pingpong_flag; + uint32_t embed_split_size; + uint32_t round_embed_split_size; + int64_t now_l1_offset; + int64_t kv_offset; + int64_t kv_offset_rope; + uint32_t sv_round_n; + uint32_t gm_split_idx; // split_idx 用于 GM 偏移计算 + uint32_t pp_n_scalar; // block_size，用于 GM 偏移计算 +}; + +// 业务函数：初始化 TP1 QK 参数 +// 对应原始 CUBE1 stage1 中 split_idx 循环体内的参数计算 +__aicore__ __attribute__((always_inline)) inline void InitTP1QKParams( + TP1Context &ctx, uint32_t split_idx, uint32_t n_idx, TP1QKParams ¶ms) +{ + uint32_t pp_n_scalar = ctx.pp_n_scalar; + uint32_t n_loop = ctx.n_loop; + uint32_t now_idx = n_idx + split_idx; + + params.l1_kv_pingpong_flag = now_idx % 2; + + // 动态计算 qk_n / qk_round_n + if (now_idx == (n_loop - 1)) { + params.qk_n = (ctx.cur_kv_seqlen - now_idx * pp_n_scalar); + } else { + params.qk_n = pp_n_scalar; + } + params.qk_round_n = RoundUp(params.qk_n); + + // embed_split_size 默认 128，idx==4 时为 64（在编排层设置） + params.embed_split_size = 128; + params.round_embed_split_size = RoundUp(params.embed_split_size); + + // now_l1_offset 默认 0，在 LoadTP1KVMain/KVRope 中更新 + params.now_l1_offset = 0; + + // block_table_id / kv_offset / kv_offset_rope 地址计算 + uint32_t block_table_id = (uint32_t)(*(block_tables_gm + + ctx.cur_batch * max_num_blocks_per_query + ctx.start_kv / block_size + now_idx)); + params.kv_offset = (int64_t)block_table_id * block_size * stride_kv; + params.kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; + + // sv_round_n 计算（用于 S→GM 的 dstStride） + uint32_t sv_n_triu = n_loop * pp_n_scalar; + uint32_t sv_n; + if (n_idx + ctx.s_block_stack > n_loop - 1) { + sv_n = ctx.cur_kv_seqlen - n_idx * pp_n_scalar; + } else { + sv_n = pp_n_scalar * ctx.s_block_stack; + } + params.sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; + + params.gm_split_idx = split_idx; + params.pp_n_scalar = pp_n_scalar; +} + +// 业务函数：TP1 Q 数据从 L1 加载到 L0A +// 复用非 TP1 的 PlatformWaitQLoadComplete / PlatformLoadQToL0A / PlatformSetQLoadComplete +// TP1 中 q_load_coeff = m，与非 TP1 一致 +__aicore__ __attribute__((always_inline)) inline void LoadTP1QDataToL0A( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformWaitQLoadComplete(embed_split_idx); + PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, m, false); + PlatformSetQLoadComplete(embed_split_idx); +} + +// 业务函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2） +// 更新 now_l1_offset 供后续 L1→L0B 使用 +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVMainFromGM( + TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 256; + PlatformLoadTP1KVMainToL1(embed_split_idx, params.qk_n, params.qk_round_n, + params.l1_kv_pingpong_flag, params.kv_offset, + embed_split_idx * 128); +} + +// 业务函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4） +// 更新 now_l1_offset 供后续 L1→L0B 使用 +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVRopeFromGM( + TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; + PlatformLoadTP1KVRopeToL1(params.qk_n, params.qk_round_n, + params.l1_kv_pingpong_flag, params.kv_offset_rope); +} + +// 业务函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVDataToL0B( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformLoadTP1KVToL0B(embed_split_idx, params.round_embed_split_size, + params.qk_round_n, params.now_l1_offset, + params.l1_kv_pingpong_flag); +} + +// 业务函数：TP1 QK MMAD 计算 +__aicore__ __attribute__((always_inline)) inline void ComputeTP1QKMMad( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformComputeTP1QKMMad(embed_split_idx, params.embed_split_size, + m, params.qk_n, params.l1_kv_pingpong_flag); +} + +// 业务函数：TP1 QK 结果拷贝到 GM（idx 4 时） +// GM 偏移: block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar +__aicore__ __attribute__((always_inline)) inline void CopyTP1QKResultToGM( + const TP1QKParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) +{ + uint64_t gm_dst_offset = (uint64_t)block_idx * TMP_SIZE_DECODER * 4 + + (uint64_t)((n_idx / 4) % 2) * TMP_SIZE_DECODER * 2 + + params.gm_split_idx * params.pp_n_scalar; + PlatformCopyTP1QKResultToGM(m, params.qk_round_n, + params.l1_kv_pingpong_flag, + gm_dst_offset, params.sv_round_n); +} + // 业务函数：将 Q Rope 数据从 GM 搬运到 L1（ND→NZ 格式转换） // INT8 场景：用 gm_to_l1 搬到独立的 l1q_rope_buf_addr_tensor // 非INT8 场景：用 DataCopy 搬到 l1q_buf_addr_tensor 的 Q 主体之后 @@ -122,6 +247,131 @@ __aicore__ __attribute__((always_inline)) inline void LoadQRopeFromGMToL1( } } +// ==================== TP1 QK 业务函数 ==================== + +// TP1 QK 参数结构体 +struct TP1QKParams { + uint32_t qk_n; + uint32_t qk_round_n; + uint32_t l1_kv_pingpong_flag; + uint32_t embed_split_size; + uint32_t round_embed_split_size; + int64_t now_l1_offset; + int64_t kv_offset; + int64_t kv_offset_rope; + uint32_t sv_round_n; + uint32_t gm_split_idx; // split_idx 用于 GM 偏移计算 + uint32_t pp_n_scalar; // block_size，用于 GM 偏移计算 +}; + +// 业务函数：初始化 TP1 QK 参数 +// 对应原始 CUBE1 stage1 中 split_idx 循环体内的参数计算 +__aicore__ __attribute__((always_inline)) inline void InitTP1QKParams( + TP1Context &ctx, uint32_t split_idx, uint32_t n_idx, TP1QKParams ¶ms) +{ + uint32_t pp_n_scalar = ctx.pp_n_scalar; + uint32_t n_loop = ctx.n_loop; + uint32_t now_idx = n_idx + split_idx; + + params.l1_kv_pingpong_flag = now_idx % 2; + + // 动态计算 qk_n / qk_round_n + if (now_idx == (n_loop - 1)) { + params.qk_n = (ctx.cur_kv_seqlen - now_idx * pp_n_scalar); + } else { + params.qk_n = pp_n_scalar; + } + params.qk_round_n = RoundUp(params.qk_n); + + // embed_split_size 默认 128，idx==4 时为 64（在编排层设置） + params.embed_split_size = 128; + params.round_embed_split_size = RoundUp(params.embed_split_size); + + // now_l1_offset 默认 0，在 LoadTP1KVMain/KVRope 中更新 + params.now_l1_offset = 0; + + // block_table_id / kv_offset / kv_offset_rope 地址计算 + uint32_t block_table_id = (uint32_t)(*(block_tables_gm + + ctx.cur_batch * max_num_blocks_per_query + ctx.start_kv / block_size + now_idx)); + params.kv_offset = (int64_t)block_table_id * block_size * stride_kv; + params.kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; + + // sv_round_n 计算（用于 S→GM 的 dstStride） + uint32_t sv_n_triu = n_loop * pp_n_scalar; + uint32_t sv_n; + if (n_idx + ctx.s_block_stack > n_loop - 1) { + sv_n = ctx.cur_kv_seqlen - n_idx * pp_n_scalar; + } else { + sv_n = pp_n_scalar * ctx.s_block_stack; + } + params.sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; + + params.gm_split_idx = split_idx; + params.pp_n_scalar = pp_n_scalar; +} + +// 业务函数：TP1 Q 数据从 L1 加载到 L0A +// 复用非 TP1 的 PlatformWaitQLoadComplete / PlatformLoadQToL0A / PlatformSetQLoadComplete +// TP1 中 q_load_coeff = m，与非 TP1 一致 +__aicore__ __attribute__((always_inline)) inline void LoadTP1QDataToL0A( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformWaitQLoadComplete(embed_split_idx); + PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, m, false); + PlatformSetQLoadComplete(embed_split_idx); +} + +// 业务函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2） +// 更新 now_l1_offset 供后续 L1→L0B 使用 +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVMainFromGM( + TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 256; + PlatformLoadTP1KVMainToL1(embed_split_idx, params.qk_n, params.qk_round_n, + params.l1_kv_pingpong_flag, params.kv_offset, + embed_split_idx * 128); +} + +// 业务函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4） +// 更新 now_l1_offset 供后续 L1→L0B 使用 +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVRopeFromGM( + TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; + PlatformLoadTP1KVRopeToL1(params.qk_n, params.qk_round_n, + params.l1_kv_pingpong_flag, params.kv_offset_rope); +} + +// 业务函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVDataToL0B( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformLoadTP1KVToL0B(embed_split_idx, params.round_embed_split_size, + params.qk_round_n, params.now_l1_offset, + params.l1_kv_pingpong_flag); +} + +// 业务函数：TP1 QK MMAD 计算 +__aicore__ __attribute__((always_inline)) inline void ComputeTP1QKMMad( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformComputeTP1QKMMad(embed_split_idx, params.embed_split_size, + m, params.qk_n, params.l1_kv_pingpong_flag); +} + +// 业务函数：TP1 QK 结果拷贝到 GM（idx 4 时） +// GM 偏移: block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar +__aicore__ __attribute__((always_inline)) inline void CopyTP1QKResultToGM( + const TP1QKParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) +{ + uint64_t gm_dst_offset = (uint64_t)block_idx * TMP_SIZE_DECODER * 4 + + (uint64_t)((n_idx / 4) % 2) * TMP_SIZE_DECODER * 2 + + params.gm_split_idx * params.pp_n_scalar; + PlatformCopyTP1QKResultToGM(m, params.qk_round_n, + params.l1_kv_pingpong_flag, + gm_dst_offset, params.sv_round_n); +} + // 业务函数：将 KV 主体数据从 GM 搬运到 L1 // ND_FORMAT 场景：ND→NZ 格式转换（gm_to_l1） // INT8 / 非INT8 NZ 场景：NZ→NZ 格式搬运（gm_to_l1） @@ -160,6 +410,131 @@ __aicore__ __attribute__((always_inline)) inline void LoadKVMainFromGMToL1( } } +// ==================== TP1 QK 业务函数 ==================== + +// TP1 QK 参数结构体 +struct TP1QKParams { + uint32_t qk_n; + uint32_t qk_round_n; + uint32_t l1_kv_pingpong_flag; + uint32_t embed_split_size; + uint32_t round_embed_split_size; + int64_t now_l1_offset; + int64_t kv_offset; + int64_t kv_offset_rope; + uint32_t sv_round_n; + uint32_t gm_split_idx; // split_idx 用于 GM 偏移计算 + uint32_t pp_n_scalar; // block_size，用于 GM 偏移计算 +}; + +// 业务函数：初始化 TP1 QK 参数 +// 对应原始 CUBE1 stage1 中 split_idx 循环体内的参数计算 +__aicore__ __attribute__((always_inline)) inline void InitTP1QKParams( + TP1Context &ctx, uint32_t split_idx, uint32_t n_idx, TP1QKParams ¶ms) +{ + uint32_t pp_n_scalar = ctx.pp_n_scalar; + uint32_t n_loop = ctx.n_loop; + uint32_t now_idx = n_idx + split_idx; + + params.l1_kv_pingpong_flag = now_idx % 2; + + // 动态计算 qk_n / qk_round_n + if (now_idx == (n_loop - 1)) { + params.qk_n = (ctx.cur_kv_seqlen - now_idx * pp_n_scalar); + } else { + params.qk_n = pp_n_scalar; + } + params.qk_round_n = RoundUp(params.qk_n); + + // embed_split_size 默认 128，idx==4 时为 64（在编排层设置） + params.embed_split_size = 128; + params.round_embed_split_size = RoundUp(params.embed_split_size); + + // now_l1_offset 默认 0，在 LoadTP1KVMain/KVRope 中更新 + params.now_l1_offset = 0; + + // block_table_id / kv_offset / kv_offset_rope 地址计算 + uint32_t block_table_id = (uint32_t)(*(block_tables_gm + + ctx.cur_batch * max_num_blocks_per_query + ctx.start_kv / block_size + now_idx)); + params.kv_offset = (int64_t)block_table_id * block_size * stride_kv; + params.kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; + + // sv_round_n 计算（用于 S→GM 的 dstStride） + uint32_t sv_n_triu = n_loop * pp_n_scalar; + uint32_t sv_n; + if (n_idx + ctx.s_block_stack > n_loop - 1) { + sv_n = ctx.cur_kv_seqlen - n_idx * pp_n_scalar; + } else { + sv_n = pp_n_scalar * ctx.s_block_stack; + } + params.sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; + + params.gm_split_idx = split_idx; + params.pp_n_scalar = pp_n_scalar; +} + +// 业务函数：TP1 Q 数据从 L1 加载到 L0A +// 复用非 TP1 的 PlatformWaitQLoadComplete / PlatformLoadQToL0A / PlatformSetQLoadComplete +// TP1 中 q_load_coeff = m，与非 TP1 一致 +__aicore__ __attribute__((always_inline)) inline void LoadTP1QDataToL0A( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformWaitQLoadComplete(embed_split_idx); + PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, m, false); + PlatformSetQLoadComplete(embed_split_idx); +} + +// 业务函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2） +// 更新 now_l1_offset 供后续 L1→L0B 使用 +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVMainFromGM( + TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 256; + PlatformLoadTP1KVMainToL1(embed_split_idx, params.qk_n, params.qk_round_n, + params.l1_kv_pingpong_flag, params.kv_offset, + embed_split_idx * 128); +} + +// 业务函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4） +// 更新 now_l1_offset 供后续 L1→L0B 使用 +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVRopeFromGM( + TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; + PlatformLoadTP1KVRopeToL1(params.qk_n, params.qk_round_n, + params.l1_kv_pingpong_flag, params.kv_offset_rope); +} + +// 业务函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVDataToL0B( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformLoadTP1KVToL0B(embed_split_idx, params.round_embed_split_size, + params.qk_round_n, params.now_l1_offset, + params.l1_kv_pingpong_flag); +} + +// 业务函数：TP1 QK MMAD 计算 +__aicore__ __attribute__((always_inline)) inline void ComputeTP1QKMMad( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformComputeTP1QKMMad(embed_split_idx, params.embed_split_size, + m, params.qk_n, params.l1_kv_pingpong_flag); +} + +// 业务函数：TP1 QK 结果拷贝到 GM（idx 4 时） +// GM 偏移: block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar +__aicore__ __attribute__((always_inline)) inline void CopyTP1QKResultToGM( + const TP1QKParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) +{ + uint64_t gm_dst_offset = (uint64_t)block_idx * TMP_SIZE_DECODER * 4 + + (uint64_t)((n_idx / 4) % 2) * TMP_SIZE_DECODER * 2 + + params.gm_split_idx * params.pp_n_scalar; + PlatformCopyTP1QKResultToGM(m, params.qk_round_n, + params.l1_kv_pingpong_flag, + gm_dst_offset, params.sv_round_n); +} + // 业务函数：将 KV Rope 数据从 GM 搬运到 L1 // ND_FORMAT 场景：ND→NZ 格式转换（gm_to_l1），d_value=64, src_d_value=stride_kv_rope // INT8 / 非INT8 NZ 场景：NZ→NZ 格式搬运（gm_to_l1），d_value=64, src_d_value=0 @@ -345,4 +720,598 @@ __aicore__ __attribute__((always_inline)) inline void ComputeQRope( params.l1_kv_pingpong_flag, n_idx, m, params.qk_round_n); } +} + +// ==================== TP1 QK 业务函数 ==================== + +// TP1 QK 参数结构体 +struct TP1QKParams { + uint32_t qk_n; + uint32_t qk_round_n; + uint32_t l1_kv_pingpong_flag; + uint32_t embed_split_size; + uint32_t round_embed_split_size; + int64_t now_l1_offset; + int64_t kv_offset; + int64_t kv_offset_rope; + uint32_t sv_round_n; + uint32_t gm_split_idx; // split_idx 用于 GM 偏移计算 + uint32_t pp_n_scalar; // block_size，用于 GM 偏移计算 +}; + +// 业务函数：初始化 TP1 QK 参数 +// 对应原始 CUBE1 stage1 中 split_idx 循环体内的参数计算 +__aicore__ __attribute__((always_inline)) inline void InitTP1QKParams( + TP1Context &ctx, uint32_t split_idx, uint32_t n_idx, TP1QKParams ¶ms) +{ + uint32_t pp_n_scalar = ctx.pp_n_scalar; + uint32_t n_loop = ctx.n_loop; + uint32_t now_idx = n_idx + split_idx; + + params.l1_kv_pingpong_flag = now_idx % 2; + + // 动态计算 qk_n / qk_round_n + if (now_idx == (n_loop - 1)) { + params.qk_n = (ctx.cur_kv_seqlen - now_idx * pp_n_scalar); + } else { + params.qk_n = pp_n_scalar; + } + params.qk_round_n = RoundUp(params.qk_n); + + // embed_split_size 默认 128，idx==4 时为 64（在编排层设置） + params.embed_split_size = 128; + params.round_embed_split_size = RoundUp(params.embed_split_size); + + // now_l1_offset 默认 0，在 LoadTP1KVMain/KVRope 中更新 + params.now_l1_offset = 0; + + // block_table_id / kv_offset / kv_offset_rope 地址计算 + uint32_t block_table_id = (uint32_t)(*(block_tables_gm + + ctx.cur_batch * max_num_blocks_per_query + ctx.start_kv / block_size + now_idx)); + params.kv_offset = (int64_t)block_table_id * block_size * stride_kv; + params.kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; + + // sv_round_n 计算（用于 S→GM 的 dstStride） + uint32_t sv_n_triu = n_loop * pp_n_scalar; + uint32_t sv_n; + if (n_idx + ctx.s_block_stack > n_loop - 1) { + sv_n = ctx.cur_kv_seqlen - n_idx * pp_n_scalar; + } else { + sv_n = pp_n_scalar * ctx.s_block_stack; + } + params.sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; + + params.gm_split_idx = split_idx; + params.pp_n_scalar = pp_n_scalar; +} + +// 业务函数：TP1 Q 数据从 L1 加载到 L0A +// 复用非 TP1 的 PlatformWaitQLoadComplete / PlatformLoadQToL0A / PlatformSetQLoadComplete +// TP1 中 q_load_coeff = m，与非 TP1 一致 +__aicore__ __attribute__((always_inline)) inline void LoadTP1QDataToL0A( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformWaitQLoadComplete(embed_split_idx); + PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, m, false); + PlatformSetQLoadComplete(embed_split_idx); +} + +// 业务函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2） +// 更新 now_l1_offset 供后续 L1→L0B 使用 +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVMainFromGM( + TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 256; + PlatformLoadTP1KVMainToL1(embed_split_idx, params.qk_n, params.qk_round_n, + params.l1_kv_pingpong_flag, params.kv_offset, + embed_split_idx * 128); +} + +// 业务函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4） +// 更新 now_l1_offset 供后续 L1→L0B 使用 +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVRopeFromGM( + TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; + PlatformLoadTP1KVRopeToL1(params.qk_n, params.qk_round_n, + params.l1_kv_pingpong_flag, params.kv_offset_rope); +} + +// 业务函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVDataToL0B( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformLoadTP1KVToL0B(embed_split_idx, params.round_embed_split_size, + params.qk_round_n, params.now_l1_offset, + params.l1_kv_pingpong_flag); +} + +// 业务函数：TP1 QK MMAD 计算 +__aicore__ __attribute__((always_inline)) inline void ComputeTP1QKMMad( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformComputeTP1QKMMad(embed_split_idx, params.embed_split_size, + m, params.qk_n, params.l1_kv_pingpong_flag); +} + +// 业务函数：TP1 QK 结果拷贝到 GM（idx 4 时） +// GM 偏移: block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar +__aicore__ __attribute__((always_inline)) inline void CopyTP1QKResultToGM( + const TP1QKParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) +{ + uint64_t gm_dst_offset = (uint64_t)block_idx * TMP_SIZE_DECODER * 4 + + (uint64_t)((n_idx / 4) % 2) * TMP_SIZE_DECODER * 2 + + params.gm_split_idx * params.pp_n_scalar; + PlatformCopyTP1QKResultToGM(m, params.qk_round_n, + params.l1_kv_pingpong_flag, + gm_dst_offset, params.sv_round_n); +} + +// 业务函数：调度 Cube MLA 任务（Run 方法的 for 循环部分） +// 遍历所有 batch，解析 tiling 参数，调用 InnerRunCubeMLA +__aicore__ __attribute__((always_inline)) inline void ScheduleCubeTasks() +{ + uint32_t q_block_num_per_batch = (q_heads + cur_qn_blk_size - 1) / cur_qn_blk_size; + uint32_t process_num = q_block_num_per_batch * num_batches; + + for (uint32_t process = block_idx; process < process_num; process += (uint32_t)block_num) { + uint32_t cur_batch = process / q_block_num_per_batch; + if (cur_batch >= num_batches) break; + + uint32_t offset_tiling = tiling_head_size + tiling_para_size * cur_batch; + uint32_t start_core_idx = (cur_batch * q_block_num_per_batch) % block_num; + + uint32_t q_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); + if (kv_seqlen == 0) { + continue; + } + uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; + + uint32_t start_head = (process % q_block_num_per_batch) * cur_qn_blk_size; + uint32_t start_kv = 0; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen; + uint32_t cur_head_num = cur_qn_blk_size; + + InnerRunCubeMLA(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, + offset_tiling); + } +} + +// ==================== TP1 QK 业务函数 ==================== + +// TP1 QK 参数结构体 +struct TP1QKParams { + uint32_t qk_n; + uint32_t qk_round_n; + uint32_t l1_kv_pingpong_flag; + uint32_t embed_split_size; + uint32_t round_embed_split_size; + int64_t now_l1_offset; + int64_t kv_offset; + int64_t kv_offset_rope; + uint32_t sv_round_n; + uint32_t gm_split_idx; // split_idx 用于 GM 偏移计算 + uint32_t pp_n_scalar; // block_size，用于 GM 偏移计算 +}; + +// 业务函数：初始化 TP1 QK 参数 +// 对应原始 CUBE1 stage1 中 split_idx 循环体内的参数计算 +__aicore__ __attribute__((always_inline)) inline void InitTP1QKParams( + TP1Context &ctx, uint32_t split_idx, uint32_t n_idx, TP1QKParams ¶ms) +{ + uint32_t pp_n_scalar = ctx.pp_n_scalar; + uint32_t n_loop = ctx.n_loop; + uint32_t now_idx = n_idx + split_idx; + + params.l1_kv_pingpong_flag = now_idx % 2; + + // 动态计算 qk_n / qk_round_n + if (now_idx == (n_loop - 1)) { + params.qk_n = (ctx.cur_kv_seqlen - now_idx * pp_n_scalar); + } else { + params.qk_n = pp_n_scalar; + } + params.qk_round_n = RoundUp(params.qk_n); + + // embed_split_size 默认 128，idx==4 时为 64（在编排层设置） + params.embed_split_size = 128; + params.round_embed_split_size = RoundUp(params.embed_split_size); + + // now_l1_offset 默认 0，在 LoadTP1KVMain/KVRope 中更新 + params.now_l1_offset = 0; + + // block_table_id / kv_offset / kv_offset_rope 地址计算 + uint32_t block_table_id = (uint32_t)(*(block_tables_gm + + ctx.cur_batch * max_num_blocks_per_query + ctx.start_kv / block_size + now_idx)); + params.kv_offset = (int64_t)block_table_id * block_size * stride_kv; + params.kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; + + // sv_round_n 计算（用于 S→GM 的 dstStride） + uint32_t sv_n_triu = n_loop * pp_n_scalar; + uint32_t sv_n; + if (n_idx + ctx.s_block_stack > n_loop - 1) { + sv_n = ctx.cur_kv_seqlen - n_idx * pp_n_scalar; + } else { + sv_n = pp_n_scalar * ctx.s_block_stack; + } + params.sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; + + params.gm_split_idx = split_idx; + params.pp_n_scalar = pp_n_scalar; +} + +// 业务函数：TP1 Q 数据从 L1 加载到 L0A +// 复用非 TP1 的 PlatformWaitQLoadComplete / PlatformLoadQToL0A / PlatformSetQLoadComplete +// TP1 中 q_load_coeff = m，与非 TP1 一致 +__aicore__ __attribute__((always_inline)) inline void LoadTP1QDataToL0A( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformWaitQLoadComplete(embed_split_idx); + PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, m, false); + PlatformSetQLoadComplete(embed_split_idx); +} + +// 业务函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2） +// 更新 now_l1_offset 供后续 L1→L0B 使用 +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVMainFromGM( + TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 256; + PlatformLoadTP1KVMainToL1(embed_split_idx, params.qk_n, params.qk_round_n, + params.l1_kv_pingpong_flag, params.kv_offset, + embed_split_idx * 128); +} + +// 业务函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4） +// 更新 now_l1_offset 供后续 L1→L0B 使用 +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVRopeFromGM( + TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; + PlatformLoadTP1KVRopeToL1(params.qk_n, params.qk_round_n, + params.l1_kv_pingpong_flag, params.kv_offset_rope); +} + +// 业务函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVDataToL0B( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformLoadTP1KVToL0B(embed_split_idx, params.round_embed_split_size, + params.qk_round_n, params.now_l1_offset, + params.l1_kv_pingpong_flag); +} + +// 业务函数：TP1 QK MMAD 计算 +__aicore__ __attribute__((always_inline)) inline void ComputeTP1QKMMad( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformComputeTP1QKMMad(embed_split_idx, params.embed_split_size, + m, params.qk_n, params.l1_kv_pingpong_flag); +} + +// 业务函数：TP1 QK 结果拷贝到 GM（idx 4 时） +// GM 偏移: block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar +__aicore__ __attribute__((always_inline)) inline void CopyTP1QKResultToGM( + const TP1QKParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) +{ + uint64_t gm_dst_offset = (uint64_t)block_idx * TMP_SIZE_DECODER * 4 + + (uint64_t)((n_idx / 4) % 2) * TMP_SIZE_DECODER * 2 + + params.gm_split_idx * params.pp_n_scalar; + PlatformCopyTP1QKResultToGM(m, params.qk_round_n, + params.l1_kv_pingpong_flag, + gm_dst_offset, params.sv_round_n); +} + +// 业务函数：调度 Cube MLA TP1 任务（RunTP1 方法的中段业务逻辑） +// 包含主循环调度 + tail 优化三分支（cores_per_seq 动态调整） +__aicore__ __attribute__((always_inline)) inline void ScheduleCubeTasksTP1() +{ + uint32_t tail = totalTaskNum % block_num; + if constexpr (EnableOptimization) { + + } else{ + tail = 0; // control whether to run tail optimization + } + uint32_t totalTaskNumRound = totalTaskNum - tail; + + + for (uint32_t process = block_idx; process < totalTaskNumRound; process += (uint32_t)block_num) { // for task + uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; + uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + + uint32_t q_seqlen = 1; + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); + if (kv_seqlen == 0) { + continue; + } + uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; + + uint32_t start_head = 0; + uint32_t start_kv = 0; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen; + uint32_t cur_head_num = q_heads; + + InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); + } + + // suppose all seqs have same length + if (tail > 0){ + uint32_t sample_kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + tiling_head_size + 2)); + bool enableExtraOptimization = true; + if (block_num % 4 == 3) { + // cannot optimize this situation due to math problem + enableExtraOptimization = false; + } + if (sample_kv_seqlen <= 2048){ + // Too Short to benefit from optimization + enableExtraOptimization = false; + } + if (!enableExtraOptimization || tail <= block_num / 2) { + // collect all metadata + uint32_t cores_per_seq = 1; + if (0 < tail && tail <= block_num / 4) {// 6 tasks left, each works with 4 cores + cores_per_seq = 4; + if (tail == 1){ + cores_per_seq = block_num; + } + else if (tail == 2){ + cores_per_seq = block_num / 2; + } + else if(tail == 3){ + cores_per_seq = block_num / 3; + } + else if(tail == 4){ + cores_per_seq = block_num / 4; + } + } + else if(block_num / 4 < tail && tail <= block_num / 3) { // 8 tasks left, each works with 3 cores + cores_per_seq = 3; + + } + else if(block_num / 3 < tail && tail <= block_num / 2) { // 12 tasks left, each works with 2 cores + cores_per_seq = 2; + } + else { + // no extra optimization for tail > 12 + cores_per_seq = 1; + } + + if(!enableExtraOptimization){ + cores_per_seq = 1; + } + + uint32_t process = totalTaskNumRound + block_idx / cores_per_seq; + if (process < totalTaskNum) { + uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; + uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + + uint32_t q_seqlen = 1; + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); + uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; + uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; + uint32_t actual_work_cores = kv_seqlen / kv_seqlen_align + (kv_seqlen % kv_seqlen_align != 0); + // cores_per_seq = actual_work_cores; + uint32_t kv_seqlen_process = 0; + if (block_idx < block_idx / cores_per_seq * cores_per_seq + actual_work_cores){ + kv_seqlen_process = (block_idx % cores_per_seq == actual_work_cores - 1) ? + (kv_seqlen - kv_seqlen_align * (actual_work_cores - 1)) : kv_seqlen_align; + } + + if (kv_seqlen > 0 && kv_seqlen_process > 0) { + uint32_t start_head = 0; + uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen_process; + uint32_t cur_head_num = q_heads; + + // no need to modify anything in cube kernel, just call the same kernel + InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); + } + } + } + else if (tail > 3 * block_num / 4){ + // no benefit for optimizing this situation + uint32_t process = totalTaskNumRound + block_idx; + if (process < totalTaskNum) { + uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; + uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + + uint32_t q_seqlen = 1; + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); + if (kv_seqlen > 0) { + uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; + + uint32_t start_head = 0; + uint32_t start_kv = 0; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen; + uint32_t cur_head_num = q_heads; + + InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); + } + } + } + else { + // 18 >= tail >= 12 + // first 12 tasks, two cores per task + { + uint32_t cores_per_seq = 2; + uint32_t process = totalTaskNumRound + block_idx / cores_per_seq; + uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; + uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + + uint32_t q_seqlen = 1; + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); + uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; + uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; + uint32_t kv_seqlen_process = (block_idx % cores_per_seq == cores_per_seq - 1) ? + (kv_seqlen - kv_seqlen_align * (cores_per_seq - 1)) : kv_seqlen_align; + + if (kv_seqlen > 0 && kv_seqlen_process > 0) { + uint32_t start_head = 0; + uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen_process; + uint32_t cur_head_num = q_heads; + + // no need to modify anything in cube kernel, just call the same kernel + InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); + } + } + { + uint32_t cores_per_seq = 4; + uint32_t process = totalTaskNumRound + block_num / 2 + block_idx / cores_per_seq; + if (process < totalTaskNum) { + uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; + uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + + uint32_t q_seqlen = 1; + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); + uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; + uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; + uint32_t kv_seqlen_process = (block_idx % cores_per_seq == cores_per_seq - 1) ? + (kv_seqlen - kv_seqlen_align * (cores_per_seq - 1)) : kv_seqlen_align; + + if (kv_seqlen > 0 && kv_seqlen_process > 0) { + uint32_t start_head = 0; + uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen_process; + uint32_t cur_head_num = q_heads; + + // no need to modify anything in cube kernel, just call the same kernel + InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); + } + } + } + } + } +} + +// ==================== TP1 QK 业务函数 ==================== + +// TP1 QK 参数结构体 +struct TP1QKParams { + uint32_t qk_n; + uint32_t qk_round_n; + uint32_t l1_kv_pingpong_flag; + uint32_t embed_split_size; + uint32_t round_embed_split_size; + int64_t now_l1_offset; + int64_t kv_offset; + int64_t kv_offset_rope; + uint32_t sv_round_n; + uint32_t gm_split_idx; // split_idx 用于 GM 偏移计算 + uint32_t pp_n_scalar; // block_size，用于 GM 偏移计算 +}; + +// 业务函数：初始化 TP1 QK 参数 +// 对应原始 CUBE1 stage1 中 split_idx 循环体内的参数计算 +__aicore__ __attribute__((always_inline)) inline void InitTP1QKParams( + TP1Context &ctx, uint32_t split_idx, uint32_t n_idx, TP1QKParams ¶ms) +{ + uint32_t pp_n_scalar = ctx.pp_n_scalar; + uint32_t n_loop = ctx.n_loop; + uint32_t now_idx = n_idx + split_idx; + + params.l1_kv_pingpong_flag = now_idx % 2; + + // 动态计算 qk_n / qk_round_n + if (now_idx == (n_loop - 1)) { + params.qk_n = (ctx.cur_kv_seqlen - now_idx * pp_n_scalar); + } else { + params.qk_n = pp_n_scalar; + } + params.qk_round_n = RoundUp(params.qk_n); + + // embed_split_size 默认 128，idx==4 时为 64（在编排层设置） + params.embed_split_size = 128; + params.round_embed_split_size = RoundUp(params.embed_split_size); + + // now_l1_offset 默认 0，在 LoadTP1KVMain/KVRope 中更新 + params.now_l1_offset = 0; + + // block_table_id / kv_offset / kv_offset_rope 地址计算 + uint32_t block_table_id = (uint32_t)(*(block_tables_gm + + ctx.cur_batch * max_num_blocks_per_query + ctx.start_kv / block_size + now_idx)); + params.kv_offset = (int64_t)block_table_id * block_size * stride_kv; + params.kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; + + // sv_round_n 计算（用于 S→GM 的 dstStride） + uint32_t sv_n_triu = n_loop * pp_n_scalar; + uint32_t sv_n; + if (n_idx + ctx.s_block_stack > n_loop - 1) { + sv_n = ctx.cur_kv_seqlen - n_idx * pp_n_scalar; + } else { + sv_n = pp_n_scalar * ctx.s_block_stack; + } + params.sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; + + params.gm_split_idx = split_idx; + params.pp_n_scalar = pp_n_scalar; +} + +// 业务函数：TP1 Q 数据从 L1 加载到 L0A +// 复用非 TP1 的 PlatformWaitQLoadComplete / PlatformLoadQToL0A / PlatformSetQLoadComplete +// TP1 中 q_load_coeff = m，与非 TP1 一致 +__aicore__ __attribute__((always_inline)) inline void LoadTP1QDataToL0A( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformWaitQLoadComplete(embed_split_idx); + PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, m, false); + PlatformSetQLoadComplete(embed_split_idx); +} + +// 业务函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2） +// 更新 now_l1_offset 供后续 L1→L0B 使用 +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVMainFromGM( + TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 256; + PlatformLoadTP1KVMainToL1(embed_split_idx, params.qk_n, params.qk_round_n, + params.l1_kv_pingpong_flag, params.kv_offset, + embed_split_idx * 128); +} + +// 业务函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4） +// 更新 now_l1_offset 供后续 L1→L0B 使用 +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVRopeFromGM( + TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; + PlatformLoadTP1KVRopeToL1(params.qk_n, params.qk_round_n, + params.l1_kv_pingpong_flag, params.kv_offset_rope); +} + +// 业务函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） +__aicore__ __attribute__((always_inline)) inline void LoadTP1KVDataToL0B( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformLoadTP1KVToL0B(embed_split_idx, params.round_embed_split_size, + params.qk_round_n, params.now_l1_offset, + params.l1_kv_pingpong_flag); +} + +// 业务函数：TP1 QK MMAD 计算 +__aicore__ __attribute__((always_inline)) inline void ComputeTP1QKMMad( + const TP1QKParams ¶ms, uint32_t embed_split_idx) +{ + PlatformComputeTP1QKMMad(embed_split_idx, params.embed_split_size, + m, params.qk_n, params.l1_kv_pingpong_flag); +} + +// 业务函数：TP1 QK 结果拷贝到 GM（idx 4 时） +// GM 偏移: block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar +__aicore__ __attribute__((always_inline)) inline void CopyTP1QKResultToGM( + const TP1QKParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) +{ + uint64_t gm_dst_offset = (uint64_t)block_idx * TMP_SIZE_DECODER * 4 + + (uint64_t)((n_idx / 4) % 2) * TMP_SIZE_DECODER * 2 + + params.gm_split_idx * params.pp_n_scalar; + PlatformCopyTP1QKResultToGM(m, params.qk_round_n, + params.l1_kv_pingpong_flag, + gm_dst_offset, params.sv_round_n); } \ No newline at end of file From e98c7519f1a6c902f2ffbb9648bd8e9d615bc886 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Tue, 11 Aug 2026 14:26:25 +0800 Subject: [PATCH 11/24] feat: youhua mla aiv run function --- .../op_kernel/multi_latent_attention.h | 590 ++++++------------ .../multi_latent_attention_aiv_arch32.h | 32 + .../op_kernel/multi_latent_attention_aiv_bs.h | 326 ++++++++++ 3 files changed, 554 insertions(+), 394 deletions(-) diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h index 1d8b7fd..10778be 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h @@ -1134,16 +1134,7 @@ class MLADecoderAiv{ __aicore__ __attribute__((always_inline)) inline void RunTP1() { - SET_FLAG(MTE3, V, EVENT_ID0); - SET_FLAG(MTE3, MTE2, EVENT_ID0); - SET_FLAG(MTE3, MTE2, EVENT_ID1); - SET_FLAG(MTE3, MTE2, EVENT_ID2); - SET_FLAG(MTE3, MTE2, EVENT_ID3); - SET_FLAG(MTE3, MTE2, EVENT_ID4); - SET_FLAG(V, MTE2, EVENT_ID4); - SET_FLAG(V, MTE2, EVENT_ID0); - SET_FLAG(MTE3, V, EVENT_ID2); - SET_FLAG(V, MTE2, EVENT_ID2); + PlatformInitVectorPipeSyncTP1(); uint32_t tail = totalTaskNum % block_num; if constexpr (EnableOptimization) { @@ -1391,19 +1382,178 @@ class MLADecoderAiv{ } } - WAIT_FLAG(MTE3, V, EVENT_ID0); - WAIT_FLAG(MTE3, MTE2, EVENT_ID0); - WAIT_FLAG(MTE3, MTE2, EVENT_ID1); - WAIT_FLAG(MTE3, MTE2, EVENT_ID2); - WAIT_FLAG(MTE3, MTE2, EVENT_ID3); - WAIT_FLAG(MTE3, MTE2, EVENT_ID4); - WAIT_FLAG(V, MTE2, EVENT_ID0); - WAIT_FLAG(V, MTE2, EVENT_ID4); - WAIT_FLAG(MTE3, V, EVENT_ID2); - WAIT_FLAG(V, MTE2, EVENT_ID2); + PlatformWaitVectorPipeSyncTP1(); } private: + // ====== AIV refactor: context struct + sub-functions ====== + + struct VectorContext { + // input params + uint32_t cur_batch; + uint32_t start_head; + uint32_t cur_nIndx; + uint32_t cur_q_seqlen; + uint32_t cur_kv_seqlen; + uint32_t cur_head_num; + uint32_t offset_tiling; + + // addresses + uint64_t addr_o_scalar; + uint64_t addr_mask_scalar; + uint32_t mask_offset; + + // loop & size + uint32_t pp_n_scalar; + uint32_t sub_n_loop; + uint32_t real_n_loop; + uint32_t n_loop; + + // QK dims + uint32_t qk_n; + uint32_t qk_round_n; + uint32_t qk_n_2; + uint32_t qk_round_n_2; + + // head split + uint32_t sub_head_num; + uint32_t sub_m; + uint32_t head_idx; + uint64_t o_offset; + + // tail info + uint32_t tail_len; + bool prev_tail_mask; + }; + + __aicore__ __attribute__((always_inline)) inline void InitVectorContext( + VectorContext &ctx, uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, + uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, + uint32_t offset_tiling) + { + ctx.cur_batch = cur_batch; + ctx.start_head = start_head; + ctx.cur_nIndx = cur_nIndx; + ctx.cur_q_seqlen = cur_q_seqlen; + ctx.cur_kv_seqlen = cur_kv_seqlen; + ctx.cur_head_num = cur_head_num; + ctx.offset_tiling = offset_tiling; + + // addr + uint32_t addr_o_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 4 + offset_tiling)); + uint32_t addr_o_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 5 + offset_tiling)); + ctx.addr_o_scalar = (uint64_t)(((uint64_t)addr_o_high32) << 32 | addr_o_loww32); + + uint32_t addr_mask_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 6 + offset_tiling)); + uint32_t addr_mask_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 7 + offset_tiling)); + ctx.addr_mask_scalar = (uint64_t)(((uint64_t)addr_mask_high32) << 32 | addr_mask_loww32); + ctx.mask_offset = ctx.addr_mask_scalar; + + // loop & size + ctx.pp_n_scalar = block_size; + ctx.sub_n_loop = ctx.pp_n_scalar / block_size; + ctx.real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; + ctx.n_loop = (cur_kv_seqlen + ctx.pp_n_scalar - 1) / ctx.pp_n_scalar; + + // QK dims + ctx.qk_n = ctx.pp_n_scalar; + ctx.qk_round_n = RoundUp(ctx.qk_n); + ctx.qk_n_2 = ctx.pp_n_scalar; + ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); + + // head split + ctx.sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; + ctx.sub_m = ctx.sub_head_num * cur_q_seqlen; + ctx.head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; + ctx.o_offset = ctx.addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; + + // tail info + ctx.tail_len = cur_kv_seqlen - (ctx.n_loop - 1) * ctx.pp_n_scalar; + ctx.prev_tail_mask = (ctx.n_loop > 1 && ctx.tail_len < cur_q_seqlen - 1); + } + + struct VectorTP1Context { + // input params + uint32_t cur_batch; + uint32_t start_head; + uint32_t cur_nIndx; + uint32_t cur_q_seqlen; + uint32_t cur_kv_seqlen; + uint32_t cur_head_num; + uint32_t offset_tiling; + + // addresses + uint64_t addr_o_scalar; + uint64_t addr_mask_scalar; + uint32_t mask_offset; + + // loop & size + uint32_t pp_n_scalar; + uint32_t sub_n_loop; + uint32_t real_n_loop; + uint32_t n_loop; + + // QK dims + uint32_t qk_n; + uint32_t qk_round_n; + uint32_t qk_n_2; + uint32_t qk_round_n_2; + + // head split + uint32_t sub_head_num; + uint32_t sub_m; + uint32_t head_idx; + uint64_t o_offset; + + // TP1 specific + uint32_t s_block_stack; + uint32_t m_slice; + uint32_t m_end; + }; + + __aicore__ __attribute__((always_inline)) inline void InitVectorTP1Context( + VectorTP1Context &ctx, uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, + uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, + uint32_t offset_tiling) + { + ctx.cur_batch = cur_batch; + ctx.start_head = start_head; + ctx.cur_nIndx = cur_nIndx; + ctx.cur_q_seqlen = cur_q_seqlen; + ctx.cur_kv_seqlen = cur_kv_seqlen; + ctx.cur_head_num = cur_head_num; + ctx.offset_tiling = offset_tiling; + + // addr + uint32_t prev_task = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); + ctx.addr_o_scalar = prev_task * q_heads * embedding_size; + ctx.addr_mask_scalar = 0; + ctx.mask_offset = 0; + + // loop & size + ctx.pp_n_scalar = block_size; + ctx.sub_n_loop = ctx.pp_n_scalar / block_size; + ctx.real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; + ctx.n_loop = (cur_kv_seqlen + ctx.pp_n_scalar - 1) / ctx.pp_n_scalar; + + // QK dims + ctx.qk_n = ctx.pp_n_scalar; + ctx.qk_round_n = RoundUp(ctx.qk_n); + ctx.qk_n_2 = ctx.pp_n_scalar; + ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); + + // head split + ctx.sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; + ctx.sub_m = ctx.sub_head_num * cur_q_seqlen; + ctx.head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; + ctx.o_offset = ctx.addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; + + // TP1 specific + ctx.s_block_stack = 4; + ctx.m_slice = FLOAT_VECTOR_SIZE / ctx.s_block_stack; + ctx.m_end = (ctx.sub_m + ctx.m_slice - 1) / ctx.m_slice; + } + // [AIV 第三层平台函数 + 第二层业务函数] 已迁移至 aiv_arch32.h / aiv_bs.h #include "multi_latent_attention_aiv_arch32.h" #include "multi_latent_attention_aiv_bs.h" @@ -2933,142 +3083,19 @@ class MLADecoderAiv{ uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) { - uint32_t addr_o_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 4 + offset_tiling)); - uint32_t addr_o_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 5 + offset_tiling)); - uint64_t addr_o_scalar = (uint64_t)(((uint64_t)addr_o_high32) << 32 | addr_o_loww32); - - uint32_t addr_mask_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 6 + offset_tiling)); - uint32_t addr_mask_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 7 + offset_tiling)); - uint64_t addr_mask_scalar = (uint64_t)(((uint64_t)addr_mask_high32) << 32 | addr_mask_loww32); - - uint32_t mask_offset = addr_mask_scalar; - - uint32_t pp_n_scalar = block_size; // 64 - uint32_t sub_n_loop = pp_n_scalar / block_size; - uint32_t real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; - - uint32_t n_loop = (cur_kv_seqlen + pp_n_scalar - 1) / pp_n_scalar; - - uint32_t qk_n = pp_n_scalar; - uint32_t qk_round_n = RoundUp(qk_n); - - uint32_t qk_n_2 = pp_n_scalar; - uint32_t qk_round_n_2 = RoundUp(qk_n_2); - - // split head num to two vectors - uint32_t sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; // 16 - uint32_t sub_m = sub_head_num * cur_q_seqlen; // 16 * 3 = 48 - - uint32_t head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; // not used - - o_offset = addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; // for NSD -> SND - - uint32_t sub_m_d128 = (sub_m + 127) / 128; // up aligned to 128 - uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; + VectorContext ctx; + InitVectorContext(ctx, cur_batch, start_head, cur_nIndx, + cur_q_seqlen, cur_kv_seqlen, cur_head_num, offset_tiling); + o_offset = ctx.o_offset; uint32_t start_kv = 0; - /* if tail length smalller than q_len - 1, then need to mask the last two tile*/ - uint32_t tail_len = cur_kv_seqlen - (n_loop - 1) * pp_n_scalar; - bool prev_tail_mask = (n_loop > 1 && tail_len < cur_q_seqlen - 1); - for (uint32_t n_idx = 0; n_idx < n_loop + 1; n_idx++) { - if (n_idx != n_loop) { - bool need_mask = false; - uint32_t mask_start_offset = 0; - if (n_idx == (n_loop - 2)) { - need_mask = prev_tail_mask; - mask_start_offset = need_mask ? (tail_len + MASK_COLUMNS - 1) * MASK_COLUMNS : 0; - } - if (n_idx == (n_loop - 1)) { - qk_n = (cur_kv_seqlen - n_idx * pp_n_scalar); - qk_round_n = RoundUp<16>(qk_n); - need_mask = true; - mask_start_offset = (qk_n - 1) * MASK_COLUMNS; - } - WaitFlagDev(QK_READY_DECODER); - /* ************ softmax1 stage1 ************* */ - WAIT_FLAG(MTE3, MTE2, EVENT_ID3); - if (sub_m > 0) { - if (mask_type == 3) { - mask_start_offset = mask_offset + n_idx * pp_n_scalar; - } - // input QK shape (sub_m, qk_round_n) - if (n_idx % 2 == 0){ - SoftmaxStage1( - p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + - (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n * T_BLOCK_OFFSET + (uint64_t)(n_idx % 2) * TMP_SIZE * T_BLOCK_OFFSET / 2], - s_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + - (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - s_rope_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + - (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mask_gm_tensor[mask_start_offset], - dm32_ubuf_tensor, ll_ubuf_tensor, pm32_ubuf_tensor, - n_idx, qk_n, qk_round_n, sub_m, 0, sub_n_loop, cur_batch, start_kv, real_n_loop, head_idx, pm_flag_scalar1, cur_q_seqlen, cur_kv_seqlen, need_mask - ); - } else { - SoftmaxStage1( - p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + - (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n * T_BLOCK_OFFSET + - TMP_SIZE * T_BLOCK_OFFSET / 2], - s_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + - (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + - TMP_SIZE_DECODER / 2], - s_rope_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + - (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + - TMP_SIZE_DECODER / 2], - mask_gm_tensor[mask_start_offset], - dm32_stage2_ubuf_tensor, ll_stage2_ubuf_tensor, pm32_ubuf_stage2_tensor, - n_idx, qk_n, qk_round_n, sub_m, 0, sub_n_loop, cur_batch, start_kv, real_n_loop, head_idx, pm_flag_scalar2, cur_q_seqlen, cur_kv_seqlen, need_mask - ); - } - } - FftsCrossCoreSync(SOFTMAX_READY_DECODER); - SET_FLAG(MTE3, MTE2, EVENT_ID3); + for (uint32_t n_idx = 0; n_idx < ctx.n_loop + 1; n_idx++) { + if (n_idx != ctx.n_loop) { + ScheduleSoftmaxStage1(ctx, n_idx, start_kv); } - /* ************ softmax2 stage1 ************* */ - - uint32_t process_row_num = 16; - uint32_t numhead_per_process = process_row_num / cur_q_seqlen; if (n_idx != 0) { - if (n_idx == n_loop) { - qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); - qk_round_n_2 = RoundUp(qk_n_2); - } - WaitFlagDev(UPDATE_READY_DECODER); - if (sub_m > 0) { - uint32_t head_loop = (sub_m + process_row_num - 1) / process_row_num; - - uint32_t head_res_row_num = 0; - uint32_t head_start_sblock_idx = 0; - uint32_t tail_res_row_num = 0; - - for (uint32_t head_loop_idx = 0; head_loop_idx < head_loop; ++head_loop_idx) { - uint32_t head_offset = head_loop_idx * process_row_num * round_v; - uint32_t cur_sub_m = head_loop_idx == (head_loop - 1) ? sub_m - head_loop_idx * process_row_num : process_row_num; // 15 or 3 - - // complete head num - head_start_sblock_idx = tail_res_row_num; - head_res_row_num = (cur_q_seqlen - tail_res_row_num) % cur_q_seqlen; - uint32_t cur_numhead_per_process = (cur_sub_m - head_res_row_num) / cur_q_seqlen; - tail_res_row_num = cur_sub_m - cur_numhead_per_process * cur_q_seqlen - head_res_row_num; - - uint32_t out_o_offset = head_loop_idx * numhead_per_process * round_v; // modified, round_v = 512 - - SoftmaxStage2MLAHeadLoop( - o_tmp_gm_tensor[(uint64_t)(block_idx * TMP_SIZE * 2 + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset + ((n_idx - 1) % 2) * TMP_SIZE)], - go_gm_tensor[(uint64_t)(block_idx * TMP_SIZE + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset)], - o_gm_tensor[(uint64_t)(o_offset + out_o_offset)], - dm32_ubuf_tensor[(uint64_t)((n_idx - 1) % 2 * 128 + head_loop_idx * process_row_num)], - ll_ubuf_tensor[(uint64_t)((n_idx - 1) % 2 * 256 + head_loop_idx * process_row_num)], - pm32_ubuf_tensor[(uint64_t)((n_idx - 1) % 2 * 128 + head_loop_idx * process_row_num)], - n_idx - 1, n_loop, qk_n_2, RoundUp(qk_round_n_2), cur_sub_m, o_offset, - head_idx + head_loop_idx * process_row_num, - pm_flag_scalar1, head_loop, head_loop_idx, cur_q_seqlen, sub_head_num, cur_head_num, - cur_numhead_per_process, - head_res_row_num, head_start_sblock_idx, tail_res_row_num); - } - } + ScheduleSoftmaxStage2(ctx, n_idx); } } } @@ -3350,129 +3377,18 @@ class MLADecoderAiv{ uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) { - uint32_t prev_task = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); - uint64_t addr_o_scalar = prev_task * q_heads * embedding_size; - uint64_t addr_mask_scalar = 0; - uint32_t mask_offset = addr_mask_scalar; - - uint32_t pp_n_scalar = block_size; // 64 - uint32_t sub_n_loop = pp_n_scalar / block_size; - uint32_t real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; - - uint32_t n_loop = (cur_kv_seqlen + pp_n_scalar - 1) / pp_n_scalar; - - uint32_t qk_n = pp_n_scalar; - uint32_t qk_round_n = RoundUp(qk_n); - - uint32_t qk_n_2 = pp_n_scalar; - uint32_t qk_round_n_2 = RoundUp(qk_n_2); - - // split head num to two vectors - uint32_t sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; // 16 - uint32_t sub_m = sub_head_num * cur_q_seqlen; // 16 * 3 = 48 - - uint32_t head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; // not used + VectorTP1Context ctx; + InitVectorTP1Context(ctx, cur_batch, start_head, cur_nIndx, + cur_q_seqlen, cur_kv_seqlen, cur_head_num, offset_tiling); - o_offset = addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; // for NSD -> SND + o_offset = ctx.o_offset; - uint32_t sub_m_d128 = (sub_m + 127) / 128; // up aligned to 128 - uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - - uint32_t start_kv = 0; - uint32_t s_block_stack = 4; - uint32_t m_slice = FLOAT_VECTOR_SIZE / s_block_stack; - uint32_t m_end = (sub_m + m_slice - 1) / m_slice; - for (uint32_t n_idx = 0; n_idx < n_loop + s_block_stack; n_idx += s_block_stack) { - if (n_idx < n_loop) { - if (n_idx + s_block_stack > n_loop - 1) { - qk_n = (cur_kv_seqlen - n_idx * pp_n_scalar); - } else { - qk_n = pp_n_scalar * s_block_stack; - } - qk_round_n = RoundUp<16>(qk_n); - if (sub_m == 0) { - WaitFlagDev(QK_READY_DECODER); - } - uint32_t pingpong_flag = 0; - for (uint32_t m_ind = 0; m_ind < m_end; m_ind++) { - uint32_t row_offset = m_ind * m_slice; - uint32_t curr_m = m_ind == m_end - 1 ? sub_m - row_offset : m_slice; - uint32_t s_ub_offset = pingpong_flag * 8192; - uint32_t p_gm_offset = (uint64_t)block_idx * TMP_SIZE * 2 + - (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + row_offset * qk_round_n + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE; - uint32_t s_gm_offset = (int64_t)block_idx * TMP_SIZE_DECODER * 4 + - (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + row_offset * qk_round_n + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE_DECODER * 2; - if (m_ind == 0) { - WaitFlagDev(QK_READY_DECODER); - } - if (curr_m == 0) { - continue; - } - OnlineSoftmaxStage1 ( - ls32_ubuf_tensor[s_ub_offset], - mask_ubuf_tensor, - mask_ubuf_tensor.template ReinterpretCast(), - lm32_ubuf_tensor[row_offset], - hm32_ubuf_tensor[row_offset], - gm32_ubuf_tensor[row_offset], - dm32_ubuf_tensor[((n_idx / s_block_stack) % 2) * UB_FLOAT_LINE_SIZE + row_offset], - ls32_ubuf_tensor[s_ub_offset], - ll_ubuf_tensor[row_offset], - gl32_ubuf_tensor[row_offset], - lp_ubuf_tensor[s_ub_offset * 2], - tv32_ubuf_tensor, - s_gm_tensor[s_gm_offset], - p_gm_tensor[p_gm_offset], - n_idx == 0, this->tor, - curr_m, qk_n, qk_round_n, pingpong_flag - ); - pingpong_flag = 1 - pingpong_flag; - } - FftsCrossCoreSync(SOFTMAX_READY_DECODER); + for (uint32_t n_idx = 0; n_idx < ctx.n_loop + ctx.s_block_stack; n_idx += ctx.s_block_stack) { + if (n_idx < ctx.n_loop) { + ScheduleOnlineSoftmaxStage1(ctx, n_idx); } - /* ************ softmax2 stage1 ************* */ - // PIPE_BARRIER(ALL); - uint32_t process_row_num = 16; - uint32_t numhead_per_process = process_row_num / cur_q_seqlen; - - if (n_idx >= s_block_stack) { - if (n_idx == n_loop) { - qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); - qk_round_n_2 = RoundUp(qk_n_2); - } - WaitFlagDev(UPDATE_READY_DECODER); - if (sub_m > 0) { - uint32_t head_loop = (sub_m + process_row_num - 1) / process_row_num; - - uint32_t head_res_row_num = 0; - uint32_t head_start_sblock_idx = 0; - uint32_t tail_res_row_num = 0; - - for (uint32_t head_loop_idx = 0; head_loop_idx < head_loop; ++head_loop_idx) { - uint32_t head_offset = head_loop_idx * process_row_num * round_v; - uint32_t cur_sub_m = head_loop_idx == (head_loop - 1) ? sub_m - head_loop_idx * process_row_num : process_row_num; // 15 or 3 - - // complete head num - head_start_sblock_idx = tail_res_row_num; - head_res_row_num = (cur_q_seqlen - tail_res_row_num) % cur_q_seqlen; - uint32_t cur_numhead_per_process = (cur_sub_m - head_res_row_num) / cur_q_seqlen; - tail_res_row_num = cur_sub_m - cur_numhead_per_process * cur_q_seqlen - head_res_row_num; - - uint32_t out_o_offset = head_loop_idx * numhead_per_process * round_v; // modified, round_v = 512 - - SoftmaxStage2MLAHeadLoopTP1( - o_tmp_gm_tensor[(uint64_t)(block_idx * TMP_SIZE * 2 + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset + ((n_idx / s_block_stack - 1) % 2) * TMP_SIZE)], - go_gm_tensor[(uint64_t)(block_idx * TMP_SIZE + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset)], - o_gm_tensor[(uint64_t)(o_offset + out_o_offset)], - dm32_ubuf_tensor[(uint64_t)((n_idx / s_block_stack - 1) % 2 * UB_FLOAT_LINE_SIZE + head_loop_idx * process_row_num)], - ll_ubuf_tensor[(uint64_t)((n_idx / s_block_stack - 1) % 2 * 256 + head_loop_idx * process_row_num)], - pm32_ubuf_tensor, - n_idx, n_loop, qk_n_2, RoundUp(qk_round_n_2), cur_sub_m, o_offset, head_idx, - pm_flag_scalar1, head_loop, head_loop_idx, cur_q_seqlen, sub_head_num, cur_head_num, cur_numhead_per_process, - head_res_row_num, head_start_sblock_idx, tail_res_row_num); - } - } + if (n_idx >= ctx.s_block_stack) { + ScheduleSoftmaxStage2TP1(ctx, n_idx); } } } @@ -3482,134 +3398,20 @@ class MLADecoderAiv{ uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) { - uint32_t prev_task = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); - uint64_t addr_o_scalar = prev_task * q_heads * embedding_size; - uint64_t addr_mask_scalar = 0; - uint32_t mask_offset = addr_mask_scalar; - - uint32_t pp_n_scalar = block_size; // 64 - uint32_t sub_n_loop = pp_n_scalar / block_size; - uint32_t real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; - - uint32_t n_loop = (cur_kv_seqlen + pp_n_scalar - 1) / pp_n_scalar; - - uint32_t qk_n = pp_n_scalar; - uint32_t qk_round_n = RoundUp(qk_n); - - uint32_t qk_n_2 = pp_n_scalar; - uint32_t qk_round_n_2 = RoundUp(qk_n_2); + VectorTP1Context ctx; + InitVectorTP1Context(ctx, 0, start_head, 0, + cur_q_seqlen, cur_kv_seqlen, cur_head_num, offset_tiling); - // split head num to two vectors - uint32_t sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; // 16 - uint32_t sub_m = sub_head_num * cur_q_seqlen; // 16 * 3 = 48 - - uint32_t head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; // not used + o_offset = ctx.o_offset; - o_offset = addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; // for NSD -> SND - - uint32_t sub_m_d128 = (sub_m + 127) / 128; // up aligned to 128 - uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - - uint32_t start_kv = 0; - uint32_t s_block_stack = 4; - uint32_t m_slice = FLOAT_VECTOR_SIZE / s_block_stack; - uint32_t m_end = (sub_m + m_slice - 1) / m_slice; - for (uint32_t n_idx = 0; n_idx < n_loop + s_block_stack; n_idx += s_block_stack) { - if (n_idx < n_loop) { - if (n_idx + s_block_stack > n_loop - 1) { - qk_n = (cur_kv_seqlen - n_idx * pp_n_scalar); - } else { - qk_n = pp_n_scalar * s_block_stack; - } - qk_round_n = RoundUp<16>(qk_n); - if (sub_m == 0) { - WaitFlagDev(QK_READY_DECODER); - } - uint32_t pingpong_flag = 0; - for (uint32_t m_ind = 0; m_ind < m_end; m_ind++) { - uint32_t row_offset = m_ind * m_slice; - uint32_t curr_m = m_ind == m_end - 1 ? sub_m - row_offset : m_slice; - uint32_t s_ub_offset = pingpong_flag * 8192; - uint32_t p_gm_offset = (uint64_t)block_idx * TMP_SIZE * 2 + - (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + row_offset * qk_round_n + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE; - uint32_t s_gm_offset = (int64_t)block_idx * TMP_SIZE_DECODER * 4 + - (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + row_offset * qk_round_n + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE_DECODER * 2; - if (m_ind == 0) { - WaitFlagDev(QK_READY_DECODER); - } - if (curr_m == 0) { - continue; - } - OnlineSoftmaxStage1 ( - ls32_ubuf_tensor[s_ub_offset], - mask_ubuf_tensor, - mask_ubuf_tensor.template ReinterpretCast(), - lm32_ubuf_tensor[row_offset], - hm32_ubuf_tensor[row_offset], - gm32_ubuf_tensor[row_offset], - dm32_ubuf_tensor[((n_idx / s_block_stack) % 2) * UB_FLOAT_LINE_SIZE + row_offset], - ls32_ubuf_tensor[s_ub_offset], - ll_ubuf_tensor[row_offset], - gl32_ubuf_tensor[row_offset], - lp_ubuf_tensor[s_ub_offset * 2], - tv32_ubuf_tensor, - s_gm_tensor[s_gm_offset], - p_gm_tensor[p_gm_offset], - n_idx == 0, this->tor, - curr_m, qk_n, qk_round_n, pingpong_flag - ); - pingpong_flag = 1 - pingpong_flag; - } - FftsCrossCoreSync(SOFTMAX_READY_DECODER); + for (uint32_t n_idx = 0; n_idx < ctx.n_loop + ctx.s_block_stack; n_idx += ctx.s_block_stack) { + if (n_idx < ctx.n_loop) { + ScheduleOnlineSoftmaxStage1(ctx, n_idx); } - /* ************ softmax2 stage1 ************* */ - // PIPE_BARRIER(ALL); - uint32_t process_row_num = 16; - uint32_t numhead_per_process = process_row_num / cur_q_seqlen; - - if (n_idx >= s_block_stack) { - if (n_idx == n_loop) { - qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); - qk_round_n_2 = RoundUp(qk_n_2); - } - WaitFlagDev(UPDATE_READY_DECODER); - if (sub_m > 0) { - uint32_t head_loop = (sub_m + process_row_num - 1) / process_row_num; - - uint32_t head_res_row_num = 0; - uint32_t head_start_sblock_idx = 0; - uint32_t tail_res_row_num = 0; - - for (uint32_t head_loop_idx = 0; head_loop_idx < head_loop; ++head_loop_idx) { - uint32_t head_offset = head_loop_idx * process_row_num * round_v; - uint32_t cur_sub_m = head_loop_idx == (head_loop - 1) ? sub_m - head_loop_idx * process_row_num : process_row_num; - - // complete head num - head_start_sblock_idx = tail_res_row_num; - head_res_row_num = (cur_q_seqlen - tail_res_row_num) % cur_q_seqlen; - uint32_t cur_numhead_per_process = (cur_sub_m - head_res_row_num) / cur_q_seqlen; - tail_res_row_num = cur_sub_m - cur_numhead_per_process * cur_q_seqlen - head_res_row_num; - - uint32_t out_o_offset = head_loop_idx * numhead_per_process * round_v; // round_v = 512 - - TailSoftmaxStage2MLAHeadLoopTP1( - o_tmp_gm_tensor[(uint64_t)(block_idx * TMP_SIZE * 2 + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset + ((n_idx / s_block_stack - 1) % 2) * TMP_SIZE)], - go_gm_tensor[(uint64_t)(block_idx * TMP_SIZE + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset)], - tmp_gm_tensor[(uint64_t)(block_idx * q_heads + start_head + sub_block_idx * cur_head_num * cur_q_seqlen / 2 + head_loop_idx * process_row_num)], - tmp_gm_tensor[(uint64_t)(block_num * q_heads + block_idx * q_heads + start_head + sub_block_idx * cur_head_num * cur_q_seqlen / 2 + head_loop_idx * process_row_num)], - dm32_ubuf_tensor[(uint64_t)((n_idx / s_block_stack - 1) % 2 * UB_FLOAT_LINE_SIZE + head_loop_idx * process_row_num)], - go32_ubuf_tensor, // no need for offset - gl32_ubuf_tensor[(uint64_t)(head_loop_idx * process_row_num)], - gm32_ubuf_tensor[(uint64_t)(head_loop_idx * process_row_num)], - n_idx, n_loop, qk_n_2, RoundUp(qk_round_n_2), cur_sub_m, o_offset, head_idx, - head_loop, head_loop_idx, cur_q_seqlen, sub_head_num, cur_head_num, cur_numhead_per_process - ); - } - } + if (n_idx >= ctx.s_block_stack) { + ScheduleTailSoftmaxStage2TP1(ctx, n_idx); } } - } __aicore__ __attribute__((always_inline)) inline void TailInnerGatherVectorTP1( diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h index 91c6ca7..f524500 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h @@ -36,4 +36,36 @@ __aicore__ __attribute__((always_inline)) inline void PlatformWaitVectorPipeSync WAIT_FLAG(V, MTE2, EVENT_ID4); WAIT_FLAG(MTE3, V, EVENT_ID2); WAIT_FLAG(V, MTE2, EVENT_ID2); +} + +// 平台函数：Vector TP1 管道同步初始化（RunTP1 方法头部 10 条 SET_FLAG） +// 与非TP1 相比多出 SET_FLAG(MTE3, MTE2, EVENT_ID1) +__aicore__ __attribute__((always_inline)) inline void PlatformInitVectorPipeSyncTP1() +{ + SET_FLAG(MTE3, V, EVENT_ID0); + SET_FLAG(MTE3, MTE2, EVENT_ID0); + SET_FLAG(MTE3, MTE2, EVENT_ID1); + SET_FLAG(MTE3, MTE2, EVENT_ID2); + SET_FLAG(MTE3, MTE2, EVENT_ID3); + SET_FLAG(MTE3, MTE2, EVENT_ID4); + SET_FLAG(V, MTE2, EVENT_ID4); + SET_FLAG(V, MTE2, EVENT_ID0); + SET_FLAG(MTE3, V, EVENT_ID2); + SET_FLAG(V, MTE2, EVENT_ID2); +} + +// 平台函数：等待 Vector TP1 管道同步完成（RunTP1 方法尾部 10 条 WAIT_FLAG） +// 与非TP1 相比多出 WAIT_FLAG(MTE3, MTE2, EVENT_ID1) +__aicore__ __attribute__((always_inline)) inline void PlatformWaitVectorPipeSyncTP1() +{ + WAIT_FLAG(MTE3, V, EVENT_ID0); + WAIT_FLAG(MTE3, MTE2, EVENT_ID0); + WAIT_FLAG(MTE3, MTE2, EVENT_ID1); + WAIT_FLAG(MTE3, MTE2, EVENT_ID2); + WAIT_FLAG(MTE3, MTE2, EVENT_ID3); + WAIT_FLAG(MTE3, MTE2, EVENT_ID4); + WAIT_FLAG(V, MTE2, EVENT_ID0); + WAIT_FLAG(V, MTE2, EVENT_ID4); + WAIT_FLAG(MTE3, V, EVENT_ID2); + WAIT_FLAG(V, MTE2, EVENT_ID2); } \ No newline at end of file diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h index 215afc0..1870f6f 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h @@ -10,6 +10,144 @@ // 本文件仅被 MLADecoderAiv 类 include，包含 Vector 侧业务调度函数。 // AIC 侧业务函数在 multi_latent_attention_bs.h 中。 +// ====== 非 TP1 路径：Phase 1 / Phase 2 业务调度 ====== + +// 业务函数：非TP1 Phase 1 — Softmax Stage1 调度 +// 包含 mask 计算、WaitFlag、SoftmaxStage1 调用（ping-pong）、FftsCrossCoreSync、SET_FLAG +__aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage1( + VectorContext &ctx, uint32_t n_idx, uint32_t start_kv) +{ + uint32_t cur_q_seqlen = ctx.cur_q_seqlen; + uint32_t cur_kv_seqlen = ctx.cur_kv_seqlen; + uint32_t cur_head_num = ctx.cur_head_num; + uint32_t cur_batch = ctx.cur_batch; + uint32_t pp_n_scalar = ctx.pp_n_scalar; + uint32_t sub_n_loop = ctx.sub_n_loop; + uint32_t real_n_loop = ctx.real_n_loop; + uint32_t n_loop = ctx.n_loop; + uint32_t sub_m = ctx.sub_m; + uint32_t head_idx = ctx.head_idx; + uint32_t mask_offset = ctx.mask_offset; + uint32_t tail_len = ctx.tail_len; + bool prev_tail_mask = ctx.prev_tail_mask; + + uint32_t qk_n = ctx.qk_n; + uint32_t qk_round_n = ctx.qk_round_n; + + bool need_mask = false; + uint32_t mask_start_offset = 0; + if (n_idx == (n_loop - 2)) { + need_mask = prev_tail_mask; + mask_start_offset = need_mask ? (tail_len + MASK_COLUMNS - 1) * MASK_COLUMNS : 0; + } + if (n_idx == (n_loop - 1)) { + qk_n = (cur_kv_seqlen - n_idx * pp_n_scalar); + qk_round_n = RoundUp<16>(qk_n); + need_mask = true; + mask_start_offset = (qk_n - 1) * MASK_COLUMNS; + } + WaitFlagDev(QK_READY_DECODER); + /* ************ softmax1 stage1 ************* */ + WAIT_FLAG(MTE3, MTE2, EVENT_ID3); + if (sub_m > 0) { + if (mask_type == 3) { + mask_start_offset = mask_offset + n_idx * pp_n_scalar; + } + // input QK shape (sub_m, qk_round_n) + if (n_idx % 2 == 0){ + SoftmaxStage1( + p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + + (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n * T_BLOCK_OFFSET + (uint64_t)(n_idx % 2) * TMP_SIZE * T_BLOCK_OFFSET / 2], + s_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + + (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + s_rope_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + + (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mask_gm_tensor[mask_start_offset], + dm32_ubuf_tensor, ll_ubuf_tensor, pm32_ubuf_tensor, + n_idx, qk_n, qk_round_n, sub_m, 0, sub_n_loop, cur_batch, start_kv, real_n_loop, head_idx, pm_flag_scalar1, cur_q_seqlen, cur_kv_seqlen, need_mask + ); + } else { + SoftmaxStage1( + p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + + (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n * T_BLOCK_OFFSET + + TMP_SIZE * T_BLOCK_OFFSET / 2], + s_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + + (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + + TMP_SIZE_DECODER / 2], + s_rope_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + + (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + + TMP_SIZE_DECODER / 2], + mask_gm_tensor[mask_start_offset], + dm32_stage2_ubuf_tensor, ll_stage2_ubuf_tensor, pm32_ubuf_stage2_tensor, + n_idx, qk_n, qk_round_n, sub_m, 0, sub_n_loop, cur_batch, start_kv, real_n_loop, head_idx, pm_flag_scalar2, cur_q_seqlen, cur_kv_seqlen, need_mask + ); + } + } + FftsCrossCoreSync(SOFTMAX_READY_DECODER); + + SET_FLAG(MTE3, MTE2, EVENT_ID3); +} + +// 业务函数：非TP1 Phase 2 — Softmax Stage2 调度 +// 包含 head 循环计算、WaitFlag、SoftmaxStage2MLAHeadLoop 调用 +__aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage2( + VectorContext &ctx, uint32_t n_idx) +{ + uint32_t cur_q_seqlen = ctx.cur_q_seqlen; + uint32_t cur_kv_seqlen = ctx.cur_kv_seqlen; + uint32_t cur_head_num = ctx.cur_head_num; + uint32_t pp_n_scalar = ctx.pp_n_scalar; + uint32_t n_loop = ctx.n_loop; + uint32_t sub_m = ctx.sub_m; + uint32_t sub_head_num = ctx.sub_head_num; + uint32_t head_idx = ctx.head_idx; + + uint32_t qk_n_2 = ctx.qk_n_2; + uint32_t qk_round_n_2 = ctx.qk_round_n_2; + + uint32_t process_row_num = 16; + uint32_t numhead_per_process = process_row_num / cur_q_seqlen; + + if (n_idx == n_loop) { + qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); + qk_round_n_2 = RoundUp(qk_n_2); + } + WaitFlagDev(UPDATE_READY_DECODER); + if (sub_m > 0) { + uint32_t head_loop = (sub_m + process_row_num - 1) / process_row_num; + + uint32_t head_res_row_num = 0; + uint32_t head_start_sblock_idx = 0; + uint32_t tail_res_row_num = 0; + + for (uint32_t head_loop_idx = 0; head_loop_idx < head_loop; ++head_loop_idx) { + uint32_t head_offset = head_loop_idx * process_row_num * round_v; + uint32_t cur_sub_m = head_loop_idx == (head_loop - 1) ? sub_m - head_loop_idx * process_row_num : process_row_num; // 15 or 3 + + // complete head num + head_start_sblock_idx = tail_res_row_num; + head_res_row_num = (cur_q_seqlen - tail_res_row_num) % cur_q_seqlen; + uint32_t cur_numhead_per_process = (cur_sub_m - head_res_row_num) / cur_q_seqlen; + tail_res_row_num = cur_sub_m - cur_numhead_per_process * cur_q_seqlen - head_res_row_num; + + uint32_t out_o_offset = head_loop_idx * numhead_per_process * round_v; // modified, round_v = 512 + + SoftmaxStage2MLAHeadLoop( + o_tmp_gm_tensor[(uint64_t)(block_idx * TMP_SIZE * 2 + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset + ((n_idx - 1) % 2) * TMP_SIZE)], + go_gm_tensor[(uint64_t)(block_idx * TMP_SIZE + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset)], + o_gm_tensor[(uint64_t)(o_offset + out_o_offset)], + dm32_ubuf_tensor[(uint64_t)((n_idx - 1) % 2 * 128 + head_loop_idx * process_row_num)], + ll_ubuf_tensor[(uint64_t)((n_idx - 1) % 2 * 256 + head_loop_idx * process_row_num)], + pm32_ubuf_tensor[(uint64_t)((n_idx - 1) % 2 * 128 + head_loop_idx * process_row_num)], + n_idx - 1, n_loop, qk_n_2, RoundUp(qk_round_n_2), cur_sub_m, o_offset, + head_idx + head_loop_idx * process_row_num, + pm_flag_scalar1, head_loop, head_loop_idx, cur_q_seqlen, sub_head_num, cur_head_num, + cur_numhead_per_process, + head_res_row_num, head_start_sblock_idx, tail_res_row_num); + } + } +} + // 业务函数：调度 Vector 非 TP1 任务（Run 方法的中段业务逻辑） // 包含 batch/head 循环调度，读取 tiling 参数后调用 InnerRunVectorChange __aicore__ __attribute__((always_inline)) inline void ScheduleVectorTasks() @@ -37,4 +175,192 @@ __aicore__ __attribute__((always_inline)) inline void ScheduleVectorTasks() InnerRunVectorChange(cur_batch, start_head, cur_nIndx, cur_q_seq_len, cur_kv_seqlen, cur_head_num, offset_tiling, 512, embed_split_loop_v_former); } +} + +// ====== TP1 路径：Phase 1 / Phase 2 业务调度 ====== + +// 业务函数：TP1 Phase 1 — OnlineSoftmax Stage1 调度（InnerRunVectorChangeTP1 / TailInnerRunVectorChangeTP1 共用） +// 包含 tail qk_n 计算、WaitFlag、OnlineSoftmaxStage1 调用（m_ind 循环 ping-pong）、FftsCrossCoreSync +__aicore__ __attribute__((always_inline)) inline void ScheduleOnlineSoftmaxStage1( + VectorTP1Context &ctx, uint32_t n_idx) +{ + uint32_t cur_q_seqlen = ctx.cur_q_seqlen; + uint32_t cur_head_num = ctx.cur_head_num; + uint32_t cur_kv_seqlen = ctx.cur_kv_seqlen; + uint32_t pp_n_scalar = ctx.pp_n_scalar; + uint32_t sub_m = ctx.sub_m; + uint32_t n_loop = ctx.n_loop; + uint32_t s_block_stack = ctx.s_block_stack; + uint32_t m_slice = ctx.m_slice; + uint32_t m_end = ctx.m_end; + + uint32_t qk_n = ctx.qk_n; + uint32_t qk_round_n = ctx.qk_round_n; + + if (n_idx + s_block_stack > n_loop - 1) { + qk_n = (cur_kv_seqlen - n_idx * pp_n_scalar); + } else { + qk_n = pp_n_scalar * s_block_stack; + } + qk_round_n = RoundUp<16>(qk_n); + if (sub_m == 0) { + WaitFlagDev(QK_READY_DECODER); + } + uint32_t pingpong_flag = 0; + for (uint32_t m_ind = 0; m_ind < m_end; m_ind++) { + uint32_t row_offset = m_ind * m_slice; + uint32_t curr_m = m_ind == m_end - 1 ? sub_m - row_offset : m_slice; + uint32_t s_ub_offset = pingpong_flag * 8192; + uint32_t p_gm_offset = (uint64_t)block_idx * TMP_SIZE * 2 + + (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + row_offset * qk_round_n + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE; + uint32_t s_gm_offset = (int64_t)block_idx * TMP_SIZE_DECODER * 4 + + (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + row_offset * qk_round_n + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE_DECODER * 2; + if (m_ind == 0) { + WaitFlagDev(QK_READY_DECODER); + } + if (curr_m == 0) { + continue; + } + OnlineSoftmaxStage1 ( + ls32_ubuf_tensor[s_ub_offset], + mask_ubuf_tensor, + mask_ubuf_tensor.template ReinterpretCast(), + lm32_ubuf_tensor[row_offset], + hm32_ubuf_tensor[row_offset], + gm32_ubuf_tensor[row_offset], + dm32_ubuf_tensor[((n_idx / s_block_stack) % 2) * UB_FLOAT_LINE_SIZE + row_offset], + ls32_ubuf_tensor[s_ub_offset], + ll_ubuf_tensor[row_offset], + gl32_ubuf_tensor[row_offset], + lp_ubuf_tensor[s_ub_offset * 2], + tv32_ubuf_tensor, + s_gm_tensor[s_gm_offset], + p_gm_tensor[p_gm_offset], + n_idx == 0, this->tor, + curr_m, qk_n, qk_round_n, pingpong_flag + ); + pingpong_flag = 1 - pingpong_flag; + } + FftsCrossCoreSync(SOFTMAX_READY_DECODER); +} + +// 业务函数：TP1 Phase 2 — SoftmaxStage2MLAHeadLoopTP1 调度（非 Tail 路径） +// 包含 head 循环计算、WaitFlag、SoftmaxStage2MLAHeadLoopTP1 调用 +__aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage2TP1( + VectorTP1Context &ctx, uint32_t n_idx) +{ + uint32_t cur_q_seqlen = ctx.cur_q_seqlen; + uint32_t cur_head_num = ctx.cur_head_num; + uint32_t cur_kv_seqlen = ctx.cur_kv_seqlen; + uint32_t pp_n_scalar = ctx.pp_n_scalar; + uint32_t sub_m = ctx.sub_m; + uint32_t sub_head_num = ctx.sub_head_num; + uint32_t n_loop = ctx.n_loop; + uint32_t head_idx = ctx.head_idx; + uint32_t s_block_stack = ctx.s_block_stack; + + uint32_t qk_n_2 = ctx.qk_n_2; + uint32_t qk_round_n_2 = ctx.qk_round_n_2; + + uint32_t process_row_num = 16; + uint32_t numhead_per_process = process_row_num / cur_q_seqlen; + + if (n_idx == n_loop) { + qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); + qk_round_n_2 = RoundUp(qk_n_2); + } + WaitFlagDev(UPDATE_READY_DECODER); + if (sub_m > 0) { + uint32_t head_loop = (sub_m + process_row_num - 1) / process_row_num; + + uint32_t head_res_row_num = 0; + uint32_t head_start_sblock_idx = 0; + uint32_t tail_res_row_num = 0; + + for (uint32_t head_loop_idx = 0; head_loop_idx < head_loop; ++head_loop_idx) { + uint32_t head_offset = head_loop_idx * process_row_num * round_v; + uint32_t cur_sub_m = head_loop_idx == (head_loop - 1) ? sub_m - head_loop_idx * process_row_num : process_row_num; + + // complete head num + head_start_sblock_idx = tail_res_row_num; + head_res_row_num = (cur_q_seqlen - tail_res_row_num) % cur_q_seqlen; + uint32_t cur_numhead_per_process = (cur_sub_m - head_res_row_num) / cur_q_seqlen; + tail_res_row_num = cur_sub_m - cur_numhead_per_process * cur_q_seqlen - head_res_row_num; + + uint32_t out_o_offset = head_loop_idx * numhead_per_process * round_v; + + SoftmaxStage2MLAHeadLoopTP1( + o_tmp_gm_tensor[(uint64_t)(block_idx * TMP_SIZE * 2 + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset + ((n_idx / s_block_stack - 1) % 2) * TMP_SIZE)], + go_gm_tensor[(uint64_t)(block_idx * TMP_SIZE + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset)], + o_gm_tensor[(uint64_t)(o_offset + out_o_offset)], + dm32_ubuf_tensor[(uint64_t)((n_idx / s_block_stack - 1) % 2 * UB_FLOAT_LINE_SIZE + head_loop_idx * process_row_num)], + ll_ubuf_tensor[(uint64_t)((n_idx / s_block_stack - 1) % 2 * 256 + head_loop_idx * process_row_num)], + pm32_ubuf_tensor, + n_idx, n_loop, qk_n_2, RoundUp(qk_round_n_2), cur_sub_m, o_offset, head_idx, + pm_flag_scalar1, head_loop, head_loop_idx, cur_q_seqlen, sub_head_num, cur_head_num, cur_numhead_per_process, + head_res_row_num, head_start_sblock_idx, tail_res_row_num); + } + } +} + +// 业务函数：TP1 Phase 2 (Tail) — TailSoftmaxStage2MLAHeadLoopTP1 调度（Tail 路径） +// 包含 head 循环计算、WaitFlag、TailSoftmaxStage2MLAHeadLoopTP1 调用 +__aicore__ __attribute__((always_inline)) inline void ScheduleTailSoftmaxStage2TP1( + VectorTP1Context &ctx, uint32_t n_idx) +{ + uint32_t cur_q_seqlen = ctx.cur_q_seqlen; + uint32_t cur_head_num = ctx.cur_head_num; + uint32_t cur_kv_seqlen = ctx.cur_kv_seqlen; + uint32_t pp_n_scalar = ctx.pp_n_scalar; + uint32_t sub_m = ctx.sub_m; + uint32_t sub_head_num = ctx.sub_head_num; + uint32_t n_loop = ctx.n_loop; + uint32_t head_idx = ctx.head_idx; + uint32_t start_head = ctx.start_head; + uint32_t s_block_stack = ctx.s_block_stack; + + uint32_t qk_n_2 = ctx.qk_n_2; + uint32_t qk_round_n_2 = ctx.qk_round_n_2; + + uint32_t process_row_num = 16; + uint32_t numhead_per_process = process_row_num / cur_q_seqlen; + + if (n_idx == n_loop) { + qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); + qk_round_n_2 = RoundUp(qk_n_2); + } + WaitFlagDev(UPDATE_READY_DECODER); + if (sub_m > 0) { + uint32_t head_loop = (sub_m + process_row_num - 1) / process_row_num; + + uint32_t head_res_row_num = 0; + uint32_t head_start_sblock_idx = 0; + uint32_t tail_res_row_num = 0; + + for (uint32_t head_loop_idx = 0; head_loop_idx < head_loop; ++head_loop_idx) { + uint32_t head_offset = head_loop_idx * process_row_num * round_v; + uint32_t cur_sub_m = head_loop_idx == (head_loop - 1) ? sub_m - head_loop_idx * process_row_num : process_row_num; + + // complete head num + head_start_sblock_idx = tail_res_row_num; + head_res_row_num = (cur_q_seqlen - tail_res_row_num) % cur_q_seqlen; + uint32_t cur_numhead_per_process = (cur_sub_m - head_res_row_num) / cur_q_seqlen; + tail_res_row_num = cur_sub_m - cur_numhead_per_process * cur_q_seqlen - head_res_row_num; + + uint32_t out_o_offset = head_loop_idx * numhead_per_process * round_v; + + TailSoftmaxStage2MLAHeadLoopTP1( + o_tmp_gm_tensor[(uint64_t)(block_idx * TMP_SIZE * 2 + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset + ((n_idx / s_block_stack - 1) % 2) * TMP_SIZE)], + go_gm_tensor[(uint64_t)(block_idx * TMP_SIZE + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset)], + tmp_gm_tensor[(uint64_t)(block_idx * q_heads + start_head + sub_block_idx * cur_head_num * cur_q_seqlen / 2 + head_loop_idx * process_row_num)], + tmp_gm_tensor[(uint64_t)(block_num * q_heads + block_idx * q_heads + start_head + sub_block_idx * cur_head_num * cur_q_seqlen / 2 + head_loop_idx * process_row_num)], + dm32_ubuf_tensor[(uint64_t)((n_idx / s_block_stack - 1) % 2 * UB_FLOAT_LINE_SIZE + head_loop_idx * process_row_num)], + go32_ubuf_tensor, + gl32_ubuf_tensor[(uint64_t)(head_loop_idx * process_row_num)], + gm32_ubuf_tensor[(uint64_t)(head_loop_idx * process_row_num)], + n_idx, n_loop, qk_n_2, RoundUp(qk_round_n_2), cur_sub_m, o_offset, head_idx, + head_loop, head_loop_idx, cur_q_seqlen, sub_head_num, cur_head_num, cur_numhead_per_process + ); + } + } } \ No newline at end of file From 933da2b6e7fa02a3de393768764ca83931fa6851 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Tue, 11 Aug 2026 15:06:33 +0800 Subject: [PATCH 12/24] feat: you hua the mla aiv ScheduleSoftmaxStage1 --- .../multi_latent_attention_aiv_arch32.h | 16 ++++++++++++++++ .../op_kernel/multi_latent_attention_aiv_bs.h | 15 +++++++++------ 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h index f524500..7f8527f 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h @@ -10,6 +10,22 @@ // 本文件仅被 MLADecoderAiv 类 include，包含 Vector 侧平台同步函数。 // AIC 侧平台函数在 multi_latent_attention_arch32.h 中。 +// 平台函数：非TP1 SoftmaxStage1 前置管道同步 +// 在 SoftmaxStage1 调用之前执行：等待 QK 就绪 + 等待 MTE3→MTE2 管道同步 +__aicore__ __attribute__((always_inline)) inline void PlatformSoftmaxStage1PreSync() +{ + WaitFlagDev(QK_READY_DECODER); + WAIT_FLAG(MTE3, MTE2, EVENT_ID3); +} + +// 平台函数：非TP1 SoftmaxStage1 后置管道同步 +// 在 SoftmaxStage1 调用之后执行：核间同步 + 设置 MTE3→MTE2 管道 flag +__aicore__ __attribute__((always_inline)) inline void PlatformSoftmaxStage1PostSync() +{ + FftsCrossCoreSync(SOFTMAX_READY_DECODER); + SET_FLAG(MTE3, MTE2, EVENT_ID3); +} + // 平台函数：Vector 非TP1 管道同步初始化（Run 方法头部 9 条 SET_FLAG） __aicore__ __attribute__((always_inline)) inline void PlatformInitVectorPipeSync() { diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h index 1870f6f..a288e81 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h @@ -13,7 +13,7 @@ // ====== 非 TP1 路径：Phase 1 / Phase 2 业务调度 ====== // 业务函数：非TP1 Phase 1 — Softmax Stage1 调度 -// 包含 mask 计算、WaitFlag、SoftmaxStage1 调用（ping-pong）、FftsCrossCoreSync、SET_FLAG +// 包含 mask 计算、平台同步调用、SoftmaxStage1 调用（ping-pong） __aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage1( VectorContext &ctx, uint32_t n_idx, uint32_t start_kv) { @@ -34,6 +34,7 @@ __aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage1( uint32_t qk_n = ctx.qk_n; uint32_t qk_round_n = ctx.qk_round_n; + // mask 边界判断 bool need_mask = false; uint32_t mask_start_offset = 0; if (n_idx == (n_loop - 2)) { @@ -46,9 +47,11 @@ __aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage1( need_mask = true; mask_start_offset = (qk_n - 1) * MASK_COLUMNS; } - WaitFlagDev(QK_READY_DECODER); - /* ************ softmax1 stage1 ************* */ - WAIT_FLAG(MTE3, MTE2, EVENT_ID3); + + // 平台同步：SoftmaxStage1 前置同步（WaitFlagDev + WAIT_FLAG） + PlatformSoftmaxStage1PreSync(); + + // SoftmaxStage1 调用（ping-pong） if (sub_m > 0) { if (mask_type == 3) { mask_start_offset = mask_offset + n_idx * pp_n_scalar; @@ -83,9 +86,9 @@ __aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage1( ); } } - FftsCrossCoreSync(SOFTMAX_READY_DECODER); - SET_FLAG(MTE3, MTE2, EVENT_ID3); + // 平台同步：SoftmaxStage1 后置同步（FftsCrossCoreSync + SET_FLAG） + PlatformSoftmaxStage1PostSync(); } // 业务函数：非TP1 Phase 2 — Softmax Stage2 调度 From d9c230c62cf1e991ac4679dfb7b864cfeb946f99 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Tue, 11 Aug 2026 15:26:56 +0800 Subject: [PATCH 13/24] feat: qian yi softmaxStage1 dao aiv_bs --- .../op_kernel/multi_latent_attention.h | 281 +---------------- .../op_kernel/multi_latent_attention_aiv_bs.h | 283 ++++++++++++++++++ 2 files changed, 284 insertions(+), 280 deletions(-) diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h index 10778be..2dcf6ec 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h @@ -1957,286 +1957,7 @@ class MLADecoderAiv{ PIPE_BARRIER(V); } - __aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( - AscendC::GlobalTensor p_gm_tensor, - AscendC::GlobalTensor s_gm_tensor, - AscendC::GlobalTensor s_rope_gm_tensor, - AscendC::GlobalTensor mask_gm_tensor, - AscendC::LocalTensor dm32_ubuf_tensor, - AscendC::LocalTensor ll_ubuf_tensor, - AscendC::LocalTensor pm32_ubuf_tensor, - uint32_t n_idx, - uint32_t qk_n, - uint32_t qk_round_n, - uint32_t sub_m, - uint32_t mask_offset, - const uint32_t sub_n_loop, - const uint32_t cur_batch, - const uint32_t start_kv, - const uint32_t real_n_loop, - const uint32_t head_idx, - const uint32_t pm_flag_scalar, - uint32_t cur_q_seqlen, - uint32_t cur_kv_seqlen, - bool need_mask - ) - { - uint32_t sub_m_d128 = (sub_m + 127) / 128; // up aligned to 128 - uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - float quantMax = (float)1 / (float)127; - WAIT_FLAG(V, MTE2, EVENT_ID2); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - DeQuantPerHeadImpl( - deq_scale_gm_tensor_q1[head_idx], - s_gm_tensor, - ls32_quant_ubuf_tensor, ls32_quant_ubuf_tensor.template ReinterpretCast(), - descale_q1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, qk_n, qk_round_n, 0, 1); - gm_to_ub( - ls32_ubuf_tensor.template ReinterpretCast(), - s_rope_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * qk_round_n / FLOAT_BLOCK_SIZE, - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - AscendC::Add(ls32_ubuf_tensor, ls32_ubuf_tensor, ls32_quant_ubuf_tensor, sub_m * qk_round_n); // float - PIPE_BARRIER(V); - } else { - gm_to_ub( - ls32_ubuf_tensor.template ReinterpretCast(), - s_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * qk_round_n / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - - // TODO add mask type condition - if (mask_type == 3) { - uint32_t aligned_mask_copy_len = RoundUp(qk_n); // 16 - uint32_t mask_dst_stride = (qk_round_n - aligned_mask_copy_len) / BLOCK_SIZE; // 0 - - AscendC::DataCopyPad( - mask_ubuf_tensor, - mask_gm_tensor, - AscendC::DataCopyExtParams( - cur_q_seqlen, - qk_n * 2, - maxKVSeqLen * 2 - qk_n * 2, - mask_dst_stride, - 0), - AscendC::DataCopyPadExtParams(false, 0, 0, 0) - ); - } else if (need_mask && mask_type == 4) { - AscendC::DataCopy( - mask_ubuf_tensor, - mask_gm_tensor, - AscendC::DataCopyParams( - cur_q_seqlen, // blockCount - qk_round_n * 2 / 32, // blockLen, 2 is sizeof(half) - MASK_COLUMNS * 2 / 32 - qk_round_n * 2 / 32, // srcStride - 0 // dstStride - ) - ); - } - - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - - if (mask_type == 3 || (need_mask && mask_type == 4)) { - AscendC::Cast( - mask32_ubuf_tensor, - mask_ubuf_tensor, - AscendC::RoundMode::CAST_NONE, - cur_q_seqlen * qk_round_n); - } - } - - for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { - muls_v(ls32_ubuf_tensor[vadd_idx * FLOAT_VECTOR_SIZE], - ls32_ubuf_tensor[vadd_idx * FLOAT_VECTOR_SIZE], - tor, - sub_m, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride - ); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - muls_v(ls32_ubuf_tensor[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - ls32_ubuf_tensor[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tor, - sub_m, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); - - if constexpr (tilingKeyType != TilingKeyType::TILING_INT8_DATA) { - if (mask_type == 3 || (need_mask && mask_type == 4)) { - uint32_t cur_compute_head_num = sub_m / cur_q_seqlen; - for (uint32_t i = 0; i < cur_compute_head_num; i++) { - Add( - ls32_ubuf_tensor[cur_q_seqlen * qk_round_n * i], - ls32_ubuf_tensor[cur_q_seqlen * qk_round_n * i], - mask32_ubuf_tensor, - cur_q_seqlen * qk_round_n - ); - } - PIPE_BARRIER(V); - } - } - - // *** lm = rowmax(ls) - ReduceMaxRepeatM(lm32_ubuf_tensor, ls32_ubuf_tensor, lp32_ubuf_tensor, sub_m, qk_n, qk_round_n); - // ReduceMaxChange(lm32_ubuf_tensor, ls32_ubuf_tensor, tv32_ubuf_tensor, round_sub_m, qk_n, qk_round_n); - if (n_idx != 0) { - // *** hm = vmax(lm, gm) - max_v(hm32_ubuf_tensor, - lm32_ubuf_tensor, - gm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); - // *** dm = gm - hm - sub_v(dm32_ubuf_tensor, - gm32_ubuf_tensor, - hm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); - } else { - // *** hm = lm - ub_to_ub( - hm32_ubuf_tensor, - lm32_ubuf_tensor, - 0, // sid - 1, // nBurst - round_sub_m / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - PIPE_BARRIER(V); - } - // *** gm = hm - ub_to_ub( - gm32_ubuf_tensor, - hm32_ubuf_tensor, - 0, // sid - 1, // nBurst - round_sub_m / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - PIPE_BARRIER(V); - // *** hm_block = expand_to_block(hm) - - // *** ls = ls - hm_block - TensorSubValueRepeatM(ls32_ubuf_tensor, ls32_ubuf_tensor, - hm32_ubuf_tensor, tv32_ubuf_tensor, - sub_m, round_sub_m, qk_n, qk_round_n); - // *** ls = exp(ls) - exp_v(ls32_ubuf_tensor, - ls32_ubuf_tensor, - (sub_m * qk_round_n + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - PIPE_BARRIER(V); - // *** lp = castfp32to16(ls) - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - sub_v(pm32_ubuf_tensor, - lm32_ubuf_tensor, - hm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); - exp_v(pm32_ubuf_tensor, - pm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - PIPE_BARRIER(V); - muls_v(pm32_ubuf_tensor, - pm32_ubuf_tensor, - quantMax, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - PIPE_BARRIER(V); - brcb_v( - tv32_ubuf_tensor.ReinterpretCast(), - pm32_ubuf_tensor.ReinterpretCast(), - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - QuantPerTokenImpl(lp_ubuf_tensor, ls32_ubuf_tensor, tv32_ubuf_tensor, sub_m, qk_n, qk_round_n, 1); - } else { - conv_v(lp_ubuf_tensor, - ls32_ubuf_tensor, - (sub_m * qk_round_n + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); - PIPE_BARRIER(V); - } - SET_FLAG(V, MTE3, EVENT_ID0); - WAIT_FLAG(V, MTE3, EVENT_ID0); - ub_to_gm( - p_gm_tensor, - lp_ubuf_tensor, - 0, // sid - 1, // nBurst - sub_m * qk_round_n * T_BLOCK_OFFSET / T_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - - // *** ll = rowsum(ls32) - ReduceSumRepeatM(ll_ubuf_tensor, ls32_ubuf_tensor, sub_m, qk_n, qk_round_n); - SET_FLAG(V, MTE2, EVENT_ID2); - PIPE_BARRIER(V); - } + // SoftmaxStage1 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） __aicore__ __attribute__((always_inline)) inline void SoftmaxStage2MLAHeadLoop( AscendC::GlobalTensor o_tmp_gm_tensor, diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h index a288e81..3025ba7 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h @@ -12,6 +12,289 @@ // ====== 非 TP1 路径：Phase 1 / Phase 2 业务调度 ====== +// 业务函数：非TP1 Phase 1 — SoftmaxStage1 核心计算（Vector 业务实现） +// 包含 DeQuant/gm_to_ub、mask 处理、muls、ReduceMax、exp、QuantPerToken/conv、ub_to_gm、ReduceSum +__aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( + AscendC::GlobalTensor p_gm_tensor, + AscendC::GlobalTensor s_gm_tensor, + AscendC::GlobalTensor s_rope_gm_tensor, + AscendC::GlobalTensor mask_gm_tensor, + AscendC::LocalTensor dm32_ubuf_tensor, + AscendC::LocalTensor ll_ubuf_tensor, + AscendC::LocalTensor pm32_ubuf_tensor, + uint32_t n_idx, + uint32_t qk_n, + uint32_t qk_round_n, + uint32_t sub_m, + uint32_t mask_offset, + const uint32_t sub_n_loop, + const uint32_t cur_batch, + const uint32_t start_kv, + const uint32_t real_n_loop, + const uint32_t head_idx, + const uint32_t pm_flag_scalar, + uint32_t cur_q_seqlen, + uint32_t cur_kv_seqlen, + bool need_mask +) +{ + uint32_t sub_m_d128 = (sub_m + 127) / 128; // up aligned to 128 + uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; + float quantMax = (float)1 / (float)127; + WAIT_FLAG(V, MTE2, EVENT_ID2); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + DeQuantPerHeadImpl( + deq_scale_gm_tensor_q1[head_idx], + s_gm_tensor, + ls32_quant_ubuf_tensor, ls32_quant_ubuf_tensor.template ReinterpretCast(), + descale_q1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, qk_n, qk_round_n, 0, 1); + gm_to_ub( + ls32_ubuf_tensor.template ReinterpretCast(), + s_rope_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * qk_round_n / FLOAT_BLOCK_SIZE, + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + AscendC::Add(ls32_ubuf_tensor, ls32_ubuf_tensor, ls32_quant_ubuf_tensor, sub_m * qk_round_n); // float + PIPE_BARRIER(V); + } else { + gm_to_ub( + ls32_ubuf_tensor.template ReinterpretCast(), + s_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * qk_round_n / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + + // TODO add mask type condition + if (mask_type == 3) { + uint32_t aligned_mask_copy_len = RoundUp(qk_n); // 16 + uint32_t mask_dst_stride = (qk_round_n - aligned_mask_copy_len) / BLOCK_SIZE; // 0 + + AscendC::DataCopyPad( + mask_ubuf_tensor, + mask_gm_tensor, + AscendC::DataCopyExtParams( + cur_q_seqlen, + qk_n * 2, + maxKVSeqLen * 2 - qk_n * 2, + mask_dst_stride, + 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0) + ); + } else if (need_mask && mask_type == 4) { + AscendC::DataCopy( + mask_ubuf_tensor, + mask_gm_tensor, + AscendC::DataCopyParams( + cur_q_seqlen, // blockCount + qk_round_n * 2 / 32, // blockLen, 2 is sizeof(half) + MASK_COLUMNS * 2 / 32 - qk_round_n * 2 / 32, // srcStride + 0 // dstStride + ) + ); + } + + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + + if (mask_type == 3 || (need_mask && mask_type == 4)) { + AscendC::Cast( + mask32_ubuf_tensor, + mask_ubuf_tensor, + AscendC::RoundMode::CAST_NONE, + cur_q_seqlen * qk_round_n); + } + } + + for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { + muls_v(ls32_ubuf_tensor[vadd_idx * FLOAT_VECTOR_SIZE], + ls32_ubuf_tensor[vadd_idx * FLOAT_VECTOR_SIZE], + tor, + sub_m, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride + ); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + muls_v(ls32_ubuf_tensor[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + ls32_ubuf_tensor[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tor, + sub_m, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + + if constexpr (tilingKeyType != TilingKeyType::TILING_INT8_DATA) { + if (mask_type == 3 || (need_mask && mask_type == 4)) { + uint32_t cur_compute_head_num = sub_m / cur_q_seqlen; + for (uint32_t i = 0; i < cur_compute_head_num; i++) { + Add( + ls32_ubuf_tensor[cur_q_seqlen * qk_round_n * i], + ls32_ubuf_tensor[cur_q_seqlen * qk_round_n * i], + mask32_ubuf_tensor, + cur_q_seqlen * qk_round_n + ); + } + PIPE_BARRIER(V); + } + } + + // *** lm = rowmax(ls) + ReduceMaxRepeatM(lm32_ubuf_tensor, ls32_ubuf_tensor, lp32_ubuf_tensor, sub_m, qk_n, qk_round_n); + // ReduceMaxChange(lm32_ubuf_tensor, ls32_ubuf_tensor, tv32_ubuf_tensor, round_sub_m, qk_n, qk_round_n); + if (n_idx != 0) { + // *** hm = vmax(lm, gm) + max_v(hm32_ubuf_tensor, + lm32_ubuf_tensor, + gm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + // *** dm = gm - hm + sub_v(dm32_ubuf_tensor, + gm32_ubuf_tensor, + hm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + } else { + // *** hm = lm + ub_to_ub( + hm32_ubuf_tensor, + lm32_ubuf_tensor, + 0, // sid + 1, // nBurst + round_sub_m / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + PIPE_BARRIER(V); + } + // *** gm = hm + ub_to_ub( + gm32_ubuf_tensor, + hm32_ubuf_tensor, + 0, // sid + 1, // nBurst + round_sub_m / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + PIPE_BARRIER(V); + // *** hm_block = expand_to_block(hm) + + // *** ls = ls - hm_block + TensorSubValueRepeatM(ls32_ubuf_tensor, ls32_ubuf_tensor, + hm32_ubuf_tensor, tv32_ubuf_tensor, + sub_m, round_sub_m, qk_n, qk_round_n); + // *** ls = exp(ls) + exp_v(ls32_ubuf_tensor, + ls32_ubuf_tensor, + (sub_m * qk_round_n + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + // *** lp = castfp32to16(ls) + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + sub_v(pm32_ubuf_tensor, + lm32_ubuf_tensor, + hm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + exp_v(pm32_ubuf_tensor, + pm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + muls_v(pm32_ubuf_tensor, + pm32_ubuf_tensor, + quantMax, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + brcb_v( + tv32_ubuf_tensor.ReinterpretCast(), + pm32_ubuf_tensor.ReinterpretCast(), + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + QuantPerTokenImpl(lp_ubuf_tensor, ls32_ubuf_tensor, tv32_ubuf_tensor, sub_m, qk_n, qk_round_n, 1); + } else { + conv_v(lp_ubuf_tensor, + ls32_ubuf_tensor, + (sub_m * qk_round_n + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + } + SET_FLAG(V, MTE3, EVENT_ID0); + WAIT_FLAG(V, MTE3, EVENT_ID0); + ub_to_gm( + p_gm_tensor, + lp_ubuf_tensor, + 0, // sid + 1, // nBurst + sub_m * qk_round_n * T_BLOCK_OFFSET / T_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + + // *** ll = rowsum(ls32) + ReduceSumRepeatM(ll_ubuf_tensor, ls32_ubuf_tensor, sub_m, qk_n, qk_round_n); + SET_FLAG(V, MTE2, EVENT_ID2); + PIPE_BARRIER(V); +} + // 业务函数：非TP1 Phase 1 — Softmax Stage1 调度 // 包含 mask 计算、平台同步调用、SoftmaxStage1 调用（ping-pong） __aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage1( From 1071524e8bdd385a01516e9b9ca8ec4ffc1ca982 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Tue, 11 Aug 2026 17:02:52 +0800 Subject: [PATCH 14/24] feat: youhua aiv InnerRunVectorChange --- .../op_kernel/multi_latent_attention.h | 64 ------ .../multi_latent_attention_aiv_arch32.h | 46 ++++ .../op_kernel/multi_latent_attention_aiv_bs.h | 208 +++++++++++++++--- 3 files changed, 222 insertions(+), 96 deletions(-) diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h index 2dcf6ec..d88b56e 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h @@ -2799,28 +2799,6 @@ class MLADecoderAiv{ PIPE_BARRIER(ALL); } - __aicore__ __attribute__((always_inline)) inline void InnerRunVectorChange( - uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, - uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, - uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) - { - VectorContext ctx; - InitVectorContext(ctx, cur_batch, start_head, cur_nIndx, - cur_q_seqlen, cur_kv_seqlen, cur_head_num, offset_tiling); - - o_offset = ctx.o_offset; - uint32_t start_kv = 0; - - for (uint32_t n_idx = 0; n_idx < ctx.n_loop + 1; n_idx++) { - if (n_idx != ctx.n_loop) { - ScheduleSoftmaxStage1(ctx, n_idx, start_kv); - } - if (n_idx != 0) { - ScheduleSoftmaxStage2(ctx, n_idx); - } - } - } - __aicore__ __attribute((always_inline)) inline void SoftmaxGatherTP1( AscendC::GlobalTensor o_gm_tensor, AscendC::GlobalTensor go_gm_tensor, @@ -3093,48 +3071,6 @@ class MLADecoderAiv{ AscendC::PipeBarrier(); } - __aicore__ __attribute__((always_inline)) inline void InnerRunVectorChangeTP1( - uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, - uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, - uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) - { - VectorTP1Context ctx; - InitVectorTP1Context(ctx, cur_batch, start_head, cur_nIndx, - cur_q_seqlen, cur_kv_seqlen, cur_head_num, offset_tiling); - - o_offset = ctx.o_offset; - - for (uint32_t n_idx = 0; n_idx < ctx.n_loop + ctx.s_block_stack; n_idx += ctx.s_block_stack) { - if (n_idx < ctx.n_loop) { - ScheduleOnlineSoftmaxStage1(ctx, n_idx); - } - if (n_idx >= ctx.s_block_stack) { - ScheduleSoftmaxStage2TP1(ctx, n_idx); - } - } - } - - __aicore__ __attribute__((always_inline)) inline void TailInnerRunVectorChangeTP1( - uint32_t start_head, - uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, - uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) - { - VectorTP1Context ctx; - InitVectorTP1Context(ctx, 0, start_head, 0, - cur_q_seqlen, cur_kv_seqlen, cur_head_num, offset_tiling); - - o_offset = ctx.o_offset; - - for (uint32_t n_idx = 0; n_idx < ctx.n_loop + ctx.s_block_stack; n_idx += ctx.s_block_stack) { - if (n_idx < ctx.n_loop) { - ScheduleOnlineSoftmaxStage1(ctx, n_idx); - } - if (n_idx >= ctx.s_block_stack) { - ScheduleTailSoftmaxStage2TP1(ctx, n_idx); - } - } - } - __aicore__ __attribute__((always_inline)) inline void TailInnerGatherVectorTP1( uint32_t start_head, uint32_t cur_q_seqlen, uint32_t cur_head_num, uint32_t start_block_idx, uint32_t cores_process, uint32_t offset_tiling) diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h index 7f8527f..8be75ff 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h @@ -54,6 +54,29 @@ __aicore__ __attribute__((always_inline)) inline void PlatformWaitVectorPipeSync WAIT_FLAG(V, MTE2, EVENT_ID2); } +// 平台函数：MTE2→V 管道同步（数据搬入后 V 侧可见） +// 用于 SoftmaxStage1 LoadQKData 段：gm_to_ub/DataCopy 完成后，通知 V 侧可读取 +__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSync() +{ + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); +} + +// 平台函数：V→MTE3 管道同步（V 计算完成后 MTE3 可搬出） +// 用于 SoftmaxStage1 QuantizeAndOutput 段：量化/转换完成后，通知 MTE3 可写入 GM +__aicore__ __attribute__((always_inline)) inline void PlatformVToMte3Sync() +{ + SET_FLAG(V, MTE3, EVENT_ID0); + WAIT_FLAG(V, MTE3, EVENT_ID0); +} + +// 平台函数：V→MTE2 管道通知（V 完成本轮计算，通知 MTE2 可搬入下轮数据） +// 用于 SoftmaxStage1 QuantizeAndOutput 段末尾：ping-pong 流水衔接 +__aicore__ __attribute__((always_inline)) inline void PlatformVToMte2Notify() +{ + SET_FLAG(V, MTE2, EVENT_ID2); +} + // 平台函数：Vector TP1 管道同步初始化（RunTP1 方法头部 10 条 SET_FLAG） // 与非TP1 相比多出 SET_FLAG(MTE3, MTE2, EVENT_ID1) __aicore__ __attribute__((always_inline)) inline void PlatformInitVectorPipeSyncTP1() @@ -84,4 +107,27 @@ __aicore__ __attribute__((always_inline)) inline void PlatformWaitVectorPipeSync WAIT_FLAG(V, MTE2, EVENT_ID4); WAIT_FLAG(MTE3, V, EVENT_ID2); WAIT_FLAG(V, MTE2, EVENT_ID2); +} + +// 平台函数：MTE2→V 管道同步（数据搬入后 V 侧可见） +// 用于 SoftmaxStage1 LoadQKData 段：gm_to_ub/DataCopy 完成后，通知 V 侧可读取 +__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSync() +{ + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); +} + +// 平台函数：V→MTE3 管道同步（V 计算完成后 MTE3 可搬出） +// 用于 SoftmaxStage1 QuantizeAndOutput 段：量化/转换完成后，通知 MTE3 可写入 GM +__aicore__ __attribute__((always_inline)) inline void PlatformVToMte3Sync() +{ + SET_FLAG(V, MTE3, EVENT_ID0); + WAIT_FLAG(V, MTE3, EVENT_ID0); +} + +// 平台函数：V→MTE2 管道通知（V 完成本轮计算，通知 MTE2 可搬入下轮数据） +// 用于 SoftmaxStage1 QuantizeAndOutput 段末尾：ping-pong 流水衔接 +__aicore__ __attribute__((always_inline)) inline void PlatformVToMte2Notify() +{ + SET_FLAG(V, MTE2, EVENT_ID2); } \ No newline at end of file diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h index 3025ba7..6b7aef0 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h @@ -12,36 +12,22 @@ // ====== 非 TP1 路径：Phase 1 / Phase 2 业务调度 ====== -// 业务函数：非TP1 Phase 1 — SoftmaxStage1 核心计算（Vector 业务实现） -// 包含 DeQuant/gm_to_ub、mask 处理、muls、ReduceMax、exp、QuantPerToken/conv、ub_to_gm、ReduceSum -__aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( - AscendC::GlobalTensor p_gm_tensor, +// 业务子函数1：QK 数据加载 +// INT8 路径：DeQuant + gm_to_ub(s_rope) + Add 融合 +// 非INT8 路径：gm_to_ub(s_gm) + mask 加载(DataCopyPad/DataCopy) + Cast +__aicore__ __attribute__((always_inline)) inline void LoadQKData( AscendC::GlobalTensor s_gm_tensor, AscendC::GlobalTensor s_rope_gm_tensor, AscendC::GlobalTensor mask_gm_tensor, - AscendC::LocalTensor dm32_ubuf_tensor, - AscendC::LocalTensor ll_ubuf_tensor, - AscendC::LocalTensor pm32_ubuf_tensor, uint32_t n_idx, uint32_t qk_n, uint32_t qk_round_n, uint32_t sub_m, - uint32_t mask_offset, - const uint32_t sub_n_loop, - const uint32_t cur_batch, - const uint32_t start_kv, - const uint32_t real_n_loop, - const uint32_t head_idx, - const uint32_t pm_flag_scalar, + const uint32_t head_idx, uint32_t cur_q_seqlen, - uint32_t cur_kv_seqlen, bool need_mask ) { - uint32_t sub_m_d128 = (sub_m + 127) / 128; // up aligned to 128 - uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - float quantMax = (float)1 / (float)127; WAIT_FLAG(V, MTE2, EVENT_ID2); if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { DeQuantPerHeadImpl( @@ -58,8 +44,7 @@ __aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( 0, // srcGap 0 // dstGap ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); + PlatformMte2ToVSync(); AscendC::Add(ls32_ubuf_tensor, ls32_ubuf_tensor, ls32_quant_ubuf_tensor, sub_m * qk_round_n); // float PIPE_BARRIER(V); } else { @@ -102,8 +87,7 @@ __aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( ); } - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); + PlatformMte2ToVSync(); if (mask_type == 3 || (need_mask && mask_type == 4)) { AscendC::Cast( @@ -113,7 +97,18 @@ __aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( cur_q_seqlen * qk_round_n); } } +} +// 业务子函数2：QK 缩放 + 非INT8 mask Add +// ls = ls × scale(tor)，非INT8 路径再叠加 mask +__aicore__ __attribute__((always_inline)) inline void ScaleAndMask( + uint32_t qk_n, + uint32_t qk_round_n, + uint32_t sub_m, + uint32_t cur_q_seqlen, + bool need_mask +) +{ for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { muls_v(ls32_ubuf_tensor[vadd_idx * FLOAT_VECTOR_SIZE], ls32_ubuf_tensor[vadd_idx * FLOAT_VECTOR_SIZE], @@ -154,10 +149,23 @@ __aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( PIPE_BARRIER(V); } } +} + +// 业务子函数3：Online Softmax 状态更新 +// lm = rowmax(ls) → hm = max(lm, gm) → dm = gm - hm → gm = hm +__aicore__ __attribute__((always_inline)) inline void UpdateSoftmaxState( + AscendC::LocalTensor dm32_ubuf_tensor, + uint32_t n_idx, + uint32_t sub_m, + uint32_t qk_n, + uint32_t qk_round_n +) +{ + uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; // *** lm = rowmax(ls) ReduceMaxRepeatM(lm32_ubuf_tensor, ls32_ubuf_tensor, lp32_ubuf_tensor, sub_m, qk_n, qk_round_n); - // ReduceMaxChange(lm32_ubuf_tensor, ls32_ubuf_tensor, tv32_ubuf_tensor, round_sub_m, qk_n, qk_round_n); if (n_idx != 0) { // *** hm = vmax(lm, gm) max_v(hm32_ubuf_tensor, @@ -209,7 +217,17 @@ __aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( 0 // dstGap ); PIPE_BARRIER(V); - // *** hm_block = expand_to_block(hm) +} + +// 业务子函数4：减法 + 指数 +// ls = ls - hm_block → ls = exp(ls) +__aicore__ __attribute__((always_inline)) inline void SubAndExp( + uint32_t sub_m, + uint32_t qk_n, + uint32_t qk_round_n +) +{ + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; // *** ls = ls - hm_block TensorSubValueRepeatM(ls32_ubuf_tensor, ls32_ubuf_tensor, @@ -225,7 +243,25 @@ __aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( 8 // srcRepeatStride ); PIPE_BARRIER(V); - // *** lp = castfp32to16(ls) +} + +// 业务子函数5：量化/转换输出 +// INT8: sub/exp/muls/brcb + QuantPerTokenImpl (FP32→INT8) +// 非INT8: conv_v (FP32→FP16) +// 共通: ub_to_gm(p_gm ← lp) + ReduceSum +__aicore__ __attribute__((always_inline)) inline void QuantizeAndOutput( + AscendC::GlobalTensor p_gm_tensor, + AscendC::LocalTensor ll_ubuf_tensor, + AscendC::LocalTensor pm32_ubuf_tensor, + uint32_t sub_m, + uint32_t qk_n, + uint32_t qk_round_n +) +{ + uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; + float quantMax = (float)1 / (float)127; + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { sub_v(pm32_ubuf_tensor, lm32_ubuf_tensor, @@ -277,8 +313,7 @@ __aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( ); PIPE_BARRIER(V); } - SET_FLAG(V, MTE3, EVENT_ID0); - WAIT_FLAG(V, MTE3, EVENT_ID0); + PlatformVToMte3Sync(); ub_to_gm( p_gm_tensor, lp_ubuf_tensor, @@ -291,10 +326,52 @@ __aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( // *** ll = rowsum(ls32) ReduceSumRepeatM(ll_ubuf_tensor, ls32_ubuf_tensor, sub_m, qk_n, qk_round_n); - SET_FLAG(V, MTE2, EVENT_ID2); + PlatformVToMte2Notify(); PIPE_BARRIER(V); } +// 业务函数：非TP1 Phase 1 — SoftmaxStage1 编排（调用5个业务子函数） +__aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( + AscendC::GlobalTensor p_gm_tensor, + AscendC::GlobalTensor s_gm_tensor, + AscendC::GlobalTensor s_rope_gm_tensor, + AscendC::GlobalTensor mask_gm_tensor, + AscendC::LocalTensor dm32_ubuf_tensor, + AscendC::LocalTensor ll_ubuf_tensor, + AscendC::LocalTensor pm32_ubuf_tensor, + uint32_t n_idx, + uint32_t qk_n, + uint32_t qk_round_n, + uint32_t sub_m, + uint32_t mask_offset, + const uint32_t sub_n_loop, + const uint32_t cur_batch, + const uint32_t start_kv, + const uint32_t real_n_loop, + const uint32_t head_idx, + const uint32_t pm_flag_scalar, + uint32_t cur_q_seqlen, + uint32_t cur_kv_seqlen, + bool need_mask +) +{ + // 段1：QK 数据加载（INT8:反量化+融合 / 非INT8:FP16+mask+Cast） + LoadQKData(s_gm_tensor, s_rope_gm_tensor, mask_gm_tensor, + n_idx, qk_n, qk_round_n, sub_m, head_idx, cur_q_seqlen, need_mask); + + // 段2：QK 缩放 + 非INT8 mask Add + ScaleAndMask(qk_n, qk_round_n, sub_m, cur_q_seqlen, need_mask); + + // 段3：Online Softmax 状态更新（ReduceMax + max/sub/ub_to_ub） + UpdateSoftmaxState(dm32_ubuf_tensor, n_idx, sub_m, qk_n, qk_round_n); + + // 段4：减法 + 指数（ls = ls - hm → ls = exp(ls)） + SubAndExp(sub_m, qk_n, qk_round_n); + + // 段5：量化/转换输出 + ub_to_gm + ReduceSum + QuantizeAndOutput(p_gm_tensor, ll_ubuf_tensor, pm32_ubuf_tensor, sub_m, qk_n, qk_round_n); +} + // 业务函数：非TP1 Phase 1 — Softmax Stage1 调度 // 包含 mask 计算、平台同步调用、SoftmaxStage1 调用（ping-pong） __aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage1( @@ -434,6 +511,30 @@ __aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage2( } } +// 业务函数：非 TP1 — InnerRunVectorChange（内层业务编排） +// 初始化 VectorContext + n_loop 软流水循环（Stage1↔Stage2 交替） +__aicore__ __attribute__((always_inline)) inline void InnerRunVectorChange( + uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, + uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, + uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) +{ + VectorContext ctx; + InitVectorContext(ctx, cur_batch, start_head, cur_nIndx, + cur_q_seqlen, cur_kv_seqlen, cur_head_num, offset_tiling); + + o_offset = ctx.o_offset; + uint32_t start_kv = 0; + + for (uint32_t n_idx = 0; n_idx < ctx.n_loop + 1; n_idx++) { + if (n_idx != ctx.n_loop) { + ScheduleSoftmaxStage1(ctx, n_idx, start_kv); + } + if (n_idx != 0) { + ScheduleSoftmaxStage2(ctx, n_idx); + } + } +} + // 业务函数：调度 Vector 非 TP1 任务（Run 方法的中段业务逻辑） // 包含 batch/head 循环调度，读取 tiling 参数后调用 InnerRunVectorChange __aicore__ __attribute__((always_inline)) inline void ScheduleVectorTasks() @@ -445,15 +546,12 @@ __aicore__ __attribute__((always_inline)) inline void ScheduleVectorTasks() cur_batch = process / q_block_num_per_batch; if (cur_batch >= num_batches) break; uint32_t offset_tiling = tiling_head_size + tiling_para_size * cur_batch; - uint32_t start_core_idx = (cur_batch * q_block_num_per_batch) % block_num; uint32_t q_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); if (kv_seqlen == 0) { continue; } - uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; uint32_t start_head = (process % q_block_num_per_batch) * cur_qn_blk_size; - uint32_t start_kv = 0; uint32_t cur_q_seq_len = q_seqlen; uint32_t cur_kv_seqlen = kv_seqlen; uint32_t cur_head_num = cur_qn_blk_size; @@ -649,4 +747,50 @@ __aicore__ __attribute__((always_inline)) inline void ScheduleTailSoftmaxStage2T ); } } +} + +// 业务函数：TP1 — InnerRunVectorChangeTP1（内层业务编排） +// 初始化 VectorTP1Context + n_loop 软流水循环（Stage1↔Stage2TP1 交替） +__aicore__ __attribute__((always_inline)) inline void InnerRunVectorChangeTP1( + uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, + uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, + uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) +{ + VectorTP1Context ctx; + InitVectorTP1Context(ctx, cur_batch, start_head, cur_nIndx, + cur_q_seqlen, cur_kv_seqlen, cur_head_num, offset_tiling); + + o_offset = ctx.o_offset; + + for (uint32_t n_idx = 0; n_idx < ctx.n_loop + ctx.s_block_stack; n_idx += ctx.s_block_stack) { + if (n_idx < ctx.n_loop) { + ScheduleOnlineSoftmaxStage1(ctx, n_idx); + } + if (n_idx >= ctx.s_block_stack) { + ScheduleSoftmaxStage2TP1(ctx, n_idx); + } + } +} + +// 业务函数：TP1 Tail — TailInnerRunVectorChangeTP1（内层业务编排） +// 初始化 VectorTP1Context + n_loop 软流水循环（Stage1↔TailStage2TP1 交替） +__aicore__ __attribute__((always_inline)) inline void TailInnerRunVectorChangeTP1( + uint32_t start_head, + uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, + uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) +{ + VectorTP1Context ctx; + InitVectorTP1Context(ctx, 0, start_head, 0, + cur_q_seqlen, cur_kv_seqlen, cur_head_num, offset_tiling); + + o_offset = ctx.o_offset; + + for (uint32_t n_idx = 0; n_idx < ctx.n_loop + ctx.s_block_stack; n_idx += ctx.s_block_stack) { + if (n_idx < ctx.n_loop) { + ScheduleOnlineSoftmaxStage1(ctx, n_idx); + } + if (n_idx >= ctx.s_block_stack) { + ScheduleTailSoftmaxStage2TP1(ctx, n_idx); + } + } } \ No newline at end of file From 6cd71eef857f56a762406970dd268fc8573475f4 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Wed, 12 Aug 2026 10:21:42 +0800 Subject: [PATCH 15/24] feat: chong gou arch32 --- .../op_kernel/multi_latent_attention.h | 5 +- .../multi_latent_attention_aiv_arch32.h | 54 +-- .../op_kernel/multi_latent_attention_aiv_bs.h | 428 ++++++++++++++---- 3 files changed, 377 insertions(+), 110 deletions(-) diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h index d88b56e..6d9e31e 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h @@ -1745,7 +1745,10 @@ class MLADecoderAiv{ PIPE_BARRIER(V); } - __aicore__ __attribute__((always_inline)) inline void TensorDivRepeatM( + // TensorDivRepeatM / TensorMulRepeatM / DeQuantPerHeadImpl / QuantPerTokenImpl + // 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） + + // SoftmaxStage1 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） const AscendC::LocalTensor& dst, const AscendC::LocalTensor& src, const AscendC::LocalTensor& src1, diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h index 8be75ff..54bc73c 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h @@ -54,27 +54,34 @@ __aicore__ __attribute__((always_inline)) inline void PlatformWaitVectorPipeSync WAIT_FLAG(V, MTE2, EVENT_ID2); } -// 平台函数：MTE2→V 管道同步（数据搬入后 V 侧可见） -// 用于 SoftmaxStage1 LoadQKData 段：gm_to_ub/DataCopy 完成后，通知 V 侧可读取 -__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSync() +// 平台函数：MTE2→V EVENT_ID0 同步（SET+WAIT） +// 用于 DeQuantPerHeadImpl：src 搬入后通知 V 管道可读 +__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSyncEvent0() { SET_FLAG(MTE2, V, EVENT_ID0); WAIT_FLAG(MTE2, V, EVENT_ID0); } -// 平台函数：V→MTE3 管道同步（V 计算完成后 MTE3 可搬出） -// 用于 SoftmaxStage1 QuantizeAndOutput 段：量化/转换完成后，通知 MTE3 可写入 GM -__aicore__ __attribute__((always_inline)) inline void PlatformVToMte3Sync() +// 平台函数：MTE2→V EVENT_ID2 同步（SET+WAIT） +// 用于 DeQuantPerHeadImpl：deScale 搬入后通知 V 管道可做 online 乘法 +__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSyncEvent2() { - SET_FLAG(V, MTE3, EVENT_ID0); - WAIT_FLAG(V, MTE3, EVENT_ID0); + SET_FLAG(MTE2, V, EVENT_ID2); + WAIT_FLAG(MTE2, V, EVENT_ID2); } -// 平台函数：V→MTE2 管道通知（V 完成本轮计算，通知 MTE2 可搬入下轮数据） -// 用于 SoftmaxStage1 QuantizeAndOutput 段末尾：ping-pong 流水衔接 -__aicore__ __attribute__((always_inline)) inline void PlatformVToMte2Notify() +// 平台函数：V→MTE2 管道等待（等待上一轮 V→MTE2 通知） +// 用于 SoftmaxStage1 LoadQKData 段开头：等待上一轮 QuantizeAndOutput 发出的 MTE2 可搬入通知 +__aicore__ __attribute__((always_inline)) inline void PlatformVToMte2Wait() { - SET_FLAG(V, MTE2, EVENT_ID2); + WAIT_FLAG(V, MTE2, EVENT_ID2); +} + +// 平台函数：V 管道同步屏障 +// 用于 SoftmaxStage1 LoadQKData 段 INT8 路径：Add 之后确保数据可见 +__aicore__ __attribute__((always_inline)) inline void PlatformVPipeBarrier() +{ + PIPE_BARRIER(V); } // 平台函数：Vector TP1 管道同步初始化（RunTP1 方法头部 10 条 SET_FLAG） @@ -109,25 +116,18 @@ __aicore__ __attribute__((always_inline)) inline void PlatformWaitVectorPipeSync WAIT_FLAG(V, MTE2, EVENT_ID2); } -// 平台函数：MTE2→V 管道同步（数据搬入后 V 侧可见） -// 用于 SoftmaxStage1 LoadQKData 段：gm_to_ub/DataCopy 完成后，通知 V 侧可读取 -__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSync() +// 平台函数：MTE2→V EVENT_ID0 同步（SET+WAIT） +// 用于 DeQuantPerHeadImpl：src 搬入后通知 V 管道可读 +__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSyncEvent0() { SET_FLAG(MTE2, V, EVENT_ID0); WAIT_FLAG(MTE2, V, EVENT_ID0); } -// 平台函数：V→MTE3 管道同步（V 计算完成后 MTE3 可搬出） -// 用于 SoftmaxStage1 QuantizeAndOutput 段：量化/转换完成后，通知 MTE3 可写入 GM -__aicore__ __attribute__((always_inline)) inline void PlatformVToMte3Sync() +// 平台函数：MTE2→V EVENT_ID2 同步（SET+WAIT） +// 用于 DeQuantPerHeadImpl：deScale 搬入后通知 V 管道可做 online 乘法 +__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSyncEvent2() { - SET_FLAG(V, MTE3, EVENT_ID0); - WAIT_FLAG(V, MTE3, EVENT_ID0); -} - -// 平台函数：V→MTE2 管道通知（V 完成本轮计算，通知 MTE2 可搬入下轮数据） -// 用于 SoftmaxStage1 QuantizeAndOutput 段末尾：ping-pong 流水衔接 -__aicore__ __attribute__((always_inline)) inline void PlatformVToMte2Notify() -{ - SET_FLAG(V, MTE2, EVENT_ID2); + SET_FLAG(MTE2, V, EVENT_ID2); + WAIT_FLAG(MTE2, V, EVENT_ID2); } \ No newline at end of file diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h index 6b7aef0..85d04ce 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h @@ -12,90 +12,97 @@ // ====== 非 TP1 路径：Phase 1 / Phase 2 业务调度 ====== -// 业务子函数1：QK 数据加载 -// INT8 路径：DeQuant + gm_to_ub(s_rope) + Add 融合 -// 非INT8 路径：gm_to_ub(s_gm) + mask 加载(DataCopyPad/DataCopy) + Cast -__aicore__ __attribute__((always_inline)) inline void LoadQKData( +// 业务子函数1a：QK 数据加载（INT8 路径） +// INT8 量化 QK 反量化 + RoPE 残差加载 + Add 融合 +// 平台调用：PlatformVToMte2Wait / PlatformMte2ToVSync / PlatformVPipeBarrier +__aicore__ __attribute__((always_inline)) inline void LoadQKDataInt8( AscendC::GlobalTensor s_gm_tensor, AscendC::GlobalTensor s_rope_gm_tensor, + uint32_t qk_n, + uint32_t qk_round_n, + uint32_t sub_m, + const uint32_t head_idx) +{ + PlatformVToMte2Wait(); + DeQuantPerHeadImpl( + deq_scale_gm_tensor_q1[head_idx], + s_gm_tensor, + ls32_quant_ubuf_tensor, ls32_quant_ubuf_tensor.template ReinterpretCast(), + descale_q1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, qk_n, qk_round_n, 0, 1); + gm_to_ub( + ls32_ubuf_tensor.template ReinterpretCast(), + s_rope_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * qk_round_n / FLOAT_BLOCK_SIZE, + 0, // srcGap + 0 // dstGap + ); + PlatformMte2ToVSync(); + AscendC::Add(ls32_ubuf_tensor, ls32_ubuf_tensor, ls32_quant_ubuf_tensor, sub_m * qk_round_n); // float + PlatformVPipeBarrier(); +} + +// 业务子函数1b：QK 数据加载（非INT8 路径） +// FP16/BF16 QK 直接加载 + 可选 mask 加载(DataCopyPad/DataCopy) + Cast +// 平台调用：PlatformVToMte2Wait / PlatformMte2ToVSync +__aicore__ __attribute__((always_inline)) inline void LoadQKDataFP16( + AscendC::GlobalTensor s_gm_tensor, AscendC::GlobalTensor mask_gm_tensor, - uint32_t n_idx, uint32_t qk_n, uint32_t qk_round_n, uint32_t sub_m, - const uint32_t head_idx, uint32_t cur_q_seqlen, - bool need_mask -) + bool need_mask) { - WAIT_FLAG(V, MTE2, EVENT_ID2); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - DeQuantPerHeadImpl( - deq_scale_gm_tensor_q1[head_idx], - s_gm_tensor, - ls32_quant_ubuf_tensor, ls32_quant_ubuf_tensor.template ReinterpretCast(), - descale_q1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, qk_n, qk_round_n, 0, 1); - gm_to_ub( - ls32_ubuf_tensor.template ReinterpretCast(), - s_rope_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * qk_round_n / FLOAT_BLOCK_SIZE, - 0, // srcGap - 0 // dstGap + PlatformVToMte2Wait(); + gm_to_ub( + ls32_ubuf_tensor.template ReinterpretCast(), + s_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * qk_round_n / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + + // mask 加载：mask_type==3 不规则 padding, mask_type==4 规则, 默认不加载 + if (mask_type == 3) { + uint32_t aligned_mask_copy_len = RoundUp(qk_n); // 16 + uint32_t mask_dst_stride = (qk_round_n - aligned_mask_copy_len) / BLOCK_SIZE; // 0 + + AscendC::DataCopyPad( + mask_ubuf_tensor, + mask_gm_tensor, + AscendC::DataCopyExtParams( + cur_q_seqlen, + qk_n * 2, + maxKVSeqLen * 2 - qk_n * 2, + mask_dst_stride, + 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0) ); - PlatformMte2ToVSync(); - AscendC::Add(ls32_ubuf_tensor, ls32_ubuf_tensor, ls32_quant_ubuf_tensor, sub_m * qk_round_n); // float - PIPE_BARRIER(V); - } else { - gm_to_ub( - ls32_ubuf_tensor.template ReinterpretCast(), - s_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * qk_round_n / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap + } else if (need_mask && mask_type == 4) { + AscendC::DataCopy( + mask_ubuf_tensor, + mask_gm_tensor, + AscendC::DataCopyParams( + cur_q_seqlen, // blockCount + qk_round_n * 2 / 32, // blockLen, 2 is sizeof(half) + MASK_COLUMNS * 2 / 32 - qk_round_n * 2 / 32, // srcStride + 0 // dstStride + ) ); + } - // TODO add mask type condition - if (mask_type == 3) { - uint32_t aligned_mask_copy_len = RoundUp(qk_n); // 16 - uint32_t mask_dst_stride = (qk_round_n - aligned_mask_copy_len) / BLOCK_SIZE; // 0 - - AscendC::DataCopyPad( - mask_ubuf_tensor, - mask_gm_tensor, - AscendC::DataCopyExtParams( - cur_q_seqlen, - qk_n * 2, - maxKVSeqLen * 2 - qk_n * 2, - mask_dst_stride, - 0), - AscendC::DataCopyPadExtParams(false, 0, 0, 0) - ); - } else if (need_mask && mask_type == 4) { - AscendC::DataCopy( - mask_ubuf_tensor, - mask_gm_tensor, - AscendC::DataCopyParams( - cur_q_seqlen, // blockCount - qk_round_n * 2 / 32, // blockLen, 2 is sizeof(half) - MASK_COLUMNS * 2 / 32 - qk_round_n * 2 / 32, // srcStride - 0 // dstStride - ) - ); - } - - PlatformMte2ToVSync(); + PlatformMte2ToVSync(); - if (mask_type == 3 || (need_mask && mask_type == 4)) { - AscendC::Cast( - mask32_ubuf_tensor, - mask_ubuf_tensor, - AscendC::RoundMode::CAST_NONE, - cur_q_seqlen * qk_round_n); - } + if (mask_type == 3 || (need_mask && mask_type == 4)) { + AscendC::Cast( + mask32_ubuf_tensor, + mask_ubuf_tensor, + AscendC::RoundMode::CAST_NONE, + cur_q_seqlen * qk_round_n); } } @@ -146,7 +153,7 @@ __aicore__ __attribute__((always_inline)) inline void ScaleAndMask( cur_q_seqlen * qk_round_n ); } - PIPE_BARRIER(V); + PlatformVPipeBarrier(); } } } @@ -179,7 +186,7 @@ __aicore__ __attribute__((always_inline)) inline void UpdateSoftmaxState( 8, // src0RepeatStride 8 // src1RepeatStride ); - PIPE_BARRIER(V); + PlatformVPipeBarrier(); // *** dm = gm - hm sub_v(dm32_ubuf_tensor, gm32_ubuf_tensor, @@ -192,7 +199,7 @@ __aicore__ __attribute__((always_inline)) inline void UpdateSoftmaxState( 8, // src0RepeatStride 8 // src1RepeatStride ); - PIPE_BARRIER(V); + PlatformVPipeBarrier(); } else { // *** hm = lm ub_to_ub( @@ -204,7 +211,7 @@ __aicore__ __attribute__((always_inline)) inline void UpdateSoftmaxState( 0, // srcGap 0 // dstGap ); - PIPE_BARRIER(V); + PlatformVPipeBarrier(); } // *** gm = hm ub_to_ub( @@ -274,7 +281,7 @@ __aicore__ __attribute__((always_inline)) inline void QuantizeAndOutput( 8, // src0RepeatStride 8 // src1RepeatStride ); - PIPE_BARRIER(V); + PlatformVPipeBarrier(); exp_v(pm32_ubuf_tensor, pm32_ubuf_tensor, sub_m_d64, // repeat @@ -283,7 +290,7 @@ __aicore__ __attribute__((always_inline)) inline void QuantizeAndOutput( 8, // dstRepeatStride 8 // srcRepeatStride ); - PIPE_BARRIER(V); + PlatformVPipeBarrier(); muls_v(pm32_ubuf_tensor, pm32_ubuf_tensor, quantMax, @@ -293,7 +300,7 @@ __aicore__ __attribute__((always_inline)) inline void QuantizeAndOutput( 8, // dstRepeatStride 8 // srcRepeatStride ); - PIPE_BARRIER(V); + PlatformVPipeBarrier(); brcb_v( tv32_ubuf_tensor.ReinterpretCast(), pm32_ubuf_tensor.ReinterpretCast(), @@ -311,7 +318,7 @@ __aicore__ __attribute__((always_inline)) inline void QuantizeAndOutput( 4, // dstRepeatStride 8 // srcRepeatStride ); - PIPE_BARRIER(V); + PlatformVPipeBarrier(); } PlatformVToMte3Sync(); ub_to_gm( @@ -356,8 +363,11 @@ __aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( ) { // 段1：QK 数据加载（INT8:反量化+融合 / 非INT8:FP16+mask+Cast） - LoadQKData(s_gm_tensor, s_rope_gm_tensor, mask_gm_tensor, - n_idx, qk_n, qk_round_n, sub_m, head_idx, cur_q_seqlen, need_mask); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + LoadQKDataInt8(s_gm_tensor, s_rope_gm_tensor, qk_n, qk_round_n, sub_m, head_idx); + } else { + LoadQKDataFP16(s_gm_tensor, mask_gm_tensor, qk_n, qk_round_n, sub_m, cur_q_seqlen, need_mask); + } // 段2：QK 缩放 + 非INT8 mask Add ScaleAndMask(qk_n, qk_round_n, sub_m, cur_q_seqlen, need_mask); @@ -793,4 +803,258 @@ __aicore__ __attribute__((always_inline)) inline void TailInnerRunVectorChangeTP ScheduleTailSoftmaxStage2TP1(ctx, n_idx); } } +} + +// ====== Tensor 逐行重复计算（Vector 业务基础函数）====== + +// Tensor 逐行除法（div_v 循环 + tail 处理） +__aicore__ __attribute__((always_inline)) inline void TensorDivRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& src1, + uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n) +{ + PIPE_BARRIER(V); + for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { + div_v(dst[vadd_idx * FLOAT_VECTOR_SIZE], + src[vadd_idx * FLOAT_VECTOR_SIZE], + src1, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + div_v(dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src1, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); +} + +// Tensor 逐行乘法（mul_v 循环 + tail 处理） +__aicore__ __attribute__((always_inline)) inline void TensorMulRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& src1, + uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t src1BlockStride +) { + PIPE_BARRIER(V); + for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { + mul_v(dst[vadd_idx * FLOAT_VECTOR_SIZE], + src[vadd_idx * FLOAT_VECTOR_SIZE], + src1, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + src1BlockStride, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + mul_v(dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src1, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + src1BlockStride, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); +} + +// ====== 量化/反量化业务函数 ====== + +// ---- DeQuantPerHeadImpl 业务子函数 ---- + +// 业务子函数1：加载 deScale + online 乘 quantScale +// 平台调用：PlatformMte2ToVSyncEvent2 / TensorMulRepeatM +__aicore__ __attribute__((always_inline)) inline void LoadDeScaleAndOnlineMul( + const AscendC::GlobalTensor& deScaleGm, + AscendC::LocalTensor deScaleUb, + AscendC::LocalTensor quantScale, + uint32_t sub_m, bool online) +{ + gm_to_ub_align(deScaleUb, + deScaleGm, + 0, // sid + 1, // nBurst + sub_m * sizeof(mmScaleType), // lenBurst + 0, // leftPaddingNum + 0, // rightPaddingNum + 0, // srcGap + 0 // dstGap + ); + if (online) { + PlatformMte2ToVSyncEvent2(); + TensorMulRepeatM(deScaleUb, deScaleUb, quantScale, 1, sub_m, RoundUp<16>(sub_m), 1); + } +} + +// 业务子函数2：加载 src(int32) + brcb 广播 deScale → tempScale +// 平台调用：PlatformMte2ToVSyncEvent0 / brcb_v / PlatformVPipeBarrier +__aicore__ __attribute__((always_inline)) inline void LoadSrcAndBrcbScale( + const AscendC::GlobalTensor& src, + AscendC::LocalTensor temp, + AscendC::LocalTensor deScaleUb, + AscendC::LocalTensor tempScale, + uint32_t sub_m, uint32_t qk_round_n, bool move_tensor) +{ + if (move_tensor) { + gm_to_ub( + temp, + src, + 0, // sid + 1, // nBurst + CeilDiv(sub_m * qk_round_n), // lenBurst + 0, // srcGap + 0 // dstGap + ); + } + PlatformMte2ToVSyncEvent0(); + brcb_v( + tempScale.template ReinterpretCast(), + deScaleUb.template ReinterpretCast(), + 1, // dstBlockStrides + 8, // dstRepeatStride + RoundUp<16>(sub_m) / FLOAT_BLOCK_SIZE // repeat + ); + PlatformVPipeBarrier(); +} + +// 业务子函数3：INT32→FP32 转换(conv_v) + 乘 tempScale +// 平台调用：conv_v / TensorMulRepeatM / PlatformVPipeBarrier +__aicore__ __attribute__((always_inline)) inline void ConvInt32ToFP32AndMul( + AscendC::LocalTensor dst, + AscendC::LocalTensor temp, + AscendC::LocalTensor tempScale, + uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n) +{ + uint32_t count = sub_m * qk_round_n; + uint32_t repeat_times = (count + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE; + if (repeat_times < 255) { + conv_v( + dst, // dst + temp, // src + repeat_times, // repeat_times + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + } else { + for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { + conv_v( + dst[vconv_idx * count / 2], // dst + temp[vconv_idx * count / 2], // src + (count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat_times + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + } + } + TensorMulRepeatM(dst, dst, tempScale, sub_m, qk_n, qk_round_n, 0); + PlatformVPipeBarrier(); +} + +// 逐 Head 反量化：INT32 → FP32 + deScale 乘法（编排层） +// 依赖：LoadDeScaleAndOnlineMul / LoadSrcAndBrcbScale / ConvInt32ToFP32AndMul +__aicore__ __attribute__((always_inline)) inline void DeQuantPerHeadImpl( + const AscendC::GlobalTensor& deScaleGm, + const AscendC::GlobalTensor& src, + AscendC::LocalTensor dst, + AscendC::LocalTensor temp, + AscendC::LocalTensor deScaleUb, + AscendC::LocalTensor tempScale, + AscendC::LocalTensor quantScale, + uint32_t sub_m, + uint32_t qk_n, + uint32_t qk_round_n, + bool online, + bool move_tensor +){ + LoadDeScaleAndOnlineMul(deScaleGm, deScaleUb, quantScale, sub_m, online); + LoadSrcAndBrcbScale(src, temp, deScaleUb, tempScale, sub_m, qk_round_n, move_tensor); + ConvInt32ToFP32AndMul(dst, temp, tempScale, sub_m, qk_n, qk_round_n); +} + +// 逐 Token 量化：FP32 → FP16 → INT8 +// 依赖：TensorDivRepeatM / TensorMulRepeatM / conv_v / Cast +__aicore__ __attribute__((always_inline)) inline void QuantPerTokenImpl( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& scale, + uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t pQuantOnline) +{ + if (pQuantOnline) { + // scr / scale + TensorDivRepeatM(dst.template ReinterpretCast(), src, scale, sub_m, qk_n, qk_round_n); + } else { + // scr * scale + TensorMulRepeatM(dst.template ReinterpretCast(), src, scale, sub_m, qk_n, qk_round_n, 0); + } + // src fp32 -> casttofp16 -> casttoint8 + uint32_t count = sub_m * qk_round_n; + uint32_t repeat_times = (count + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE; + if (repeat_times < 255) { + conv_v( + dst.template ReinterpretCast(), // dst + dst.template ReinterpretCast(), // src + repeat_times, // repeat_times + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + } else { + for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { + conv_v( + dst.template ReinterpretCast()[vconv_idx * count / 2], // dst + dst.template ReinterpretCast()[vconv_idx * count / 2], // src + (count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat_times + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + } + } + PIPE_BARRIER(V); + for (uint32_t row_idx = 0; row_idx < qk_n / HALF_VECTOR_SIZE; ++row_idx) { + AscendC::Cast(dst.template ReinterpretCast()[row_idx * HALF_VECTOR_SIZE], + dst.template ReinterpretCast()[row_idx * HALF_VECTOR_SIZE], AscendC::RoundMode::CAST_RINT, + (uint64_t)0, sub_m, {1, 1, (uint8_t)((qk_round_n) / BLOCK_SIZE), (uint8_t)(qk_round_n / BLOCK_SIZE)}); + } + if (qk_n % HALF_VECTOR_SIZE > 0) { + __set_mask(qk_n % HALF_VECTOR_SIZE); + AscendC::Cast(dst.template ReinterpretCast()[qk_n / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], + dst.template ReinterpretCast()[qk_n / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], AscendC::RoundMode::CAST_RINT, + (uint64_t)0, sub_m, {1, 1, (uint8_t)((qk_round_n) / BLOCK_SIZE), (uint8_t)(qk_round_n / BLOCK_SIZE)}); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); } \ No newline at end of file From 68b59f0cbf8c5208df272c7869fd79fb18390364 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Wed, 12 Aug 2026 15:28:55 +0800 Subject: [PATCH 16/24] feat: modify the code adapte the 3 ceng jiagou --- .../op_kernel/multi_latent_attention.h | 1170 +--------------- .../op_kernel/multi_latent_attention_aiv_bs.h | 1205 +++++++++++++++++ 2 files changed, 1210 insertions(+), 1165 deletions(-) diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h index 6d9e31e..6a5465a 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h @@ -1633,1174 +1633,14 @@ class MLADecoderAiv{ } - __aicore__ __attribute__((always_inline)) inline void ReduceSumRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - uint32_t sub_m, - uint32_t qk_n, - uint32_t qk_round_n) - { - if (qk_n <= FLOAT_VECTOR_SIZE) { - __set_mask(qk_n); - cadd_v( - dst, - src, - sub_m, // repeat - 1, // dstRepeatStride - 1, // srcBlockStride - qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } else { - for (uint32_t rowsum_idx = 1; rowsum_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowsum_idx) { - add_v( - src, - src, - src[rowsum_idx * FLOAT_VECTOR_SIZE], - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride - ); - PIPE_BARRIER(V); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - add_v( - src, - src, - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); - - cadd_v( - dst, - src, - sub_m, // repeat - 1, // dstRepeatStride - 1, // srcBlockStride - qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride - ); - } - } - - __aicore__ __attribute__((always_inline)) inline void TensorSubValueRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& MaxTensor, - const AscendC::LocalTensor& tempMaxTensor, - uint32_t sub_m, - uint32_t round_sub_m, - uint32_t qk_n, - uint32_t qk_round_n) - { - brcb_v( - tempMaxTensor.ReinterpretCast(), - MaxTensor.ReinterpretCast(), - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - PIPE_BARRIER(V); - for (uint32_t sub_v_idx = 0; sub_v_idx < qk_n / FLOAT_VECTOR_SIZE; ++sub_v_idx) { - sub_v(dst[sub_v_idx * FLOAT_VECTOR_SIZE], - src[sub_v_idx * FLOAT_VECTOR_SIZE], - tempMaxTensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - sub_v(dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tempMaxTensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); - } - - // TensorDivRepeatM / TensorMulRepeatM / DeQuantPerHeadImpl / QuantPerTokenImpl + // ReduceSumRepeatM // 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） - // SoftmaxStage1 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& src1, - uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n) - { - PIPE_BARRIER(V); - for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { - div_v(dst[vadd_idx * FLOAT_VECTOR_SIZE], - src[vadd_idx * FLOAT_VECTOR_SIZE], - src1, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - div_v(dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src1, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); - } - - __aicore__ __attribute__((always_inline)) inline void TensorMulRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& src1, - uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t src1BlockStride - ) { - PIPE_BARRIER(V); - for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { - mul_v(dst[vadd_idx * FLOAT_VECTOR_SIZE], - src[vadd_idx * FLOAT_VECTOR_SIZE], - src1, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - src1BlockStride, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - mul_v(dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src1, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - src1BlockStride, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); - } - - __aicore__ __attribute__((always_inline)) inline void DeQuantPerHeadImpl( - const AscendC::GlobalTensor& deScaleGm, - const AscendC::GlobalTensor& src, - AscendC::LocalTensor dst, - AscendC::LocalTensor temp, - AscendC::LocalTensor deScaleUb, - AscendC::LocalTensor tempScale, - AscendC::LocalTensor quantScale, - uint32_t sub_m, - uint32_t qk_n, - uint32_t qk_round_n, - bool online, - bool move_tensor - ){ - gm_to_ub_align(deScaleUb, - deScaleGm, - 0, // sid - 1, // nBurst - sub_m * sizeof(mmScaleType), // lenBurst - 0, // leftPaddingNum - 0, // rightPaddingNum - 0, // srcGap - 0 // dstGap - ); - if (online) { - // if dequant online need mul p quant scale - SET_FLAG(MTE2, V, EVENT_ID2); - WAIT_FLAG(MTE2, V, EVENT_ID2); - TensorMulRepeatM(deScaleUb, deScaleUb, quantScale, 1, sub_m, RoundUp<16>(sub_m), 1); - } - - if (move_tensor) { - gm_to_ub( - temp, - src, - 0, // sid - 1, // nBurst - CeilDiv(sub_m * qk_round_n), // lenBurst - 0, // srcGap - 0 // dstGap - ); - } - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - brcb_v( - tempScale.template ReinterpretCast(), - deScaleUb.template ReinterpretCast(), - 1, // dstBlockStrides - 8, // dstRepeatStride - RoundUp<16>(sub_m) / FLOAT_BLOCK_SIZE // repeat - ); - PIPE_BARRIER(V); - uint32_t count = sub_m * qk_round_n; - uint32_t repeat_times = (count + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE; - if (repeat_times < 255) { - conv_v( - dst, // dst - temp, // src - repeat_times, // repeat_times - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - } else { - for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { - conv_v( - dst[vconv_idx * count / 2], // dst - temp[vconv_idx * count / 2], // src - (count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat_times - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - } - } - TensorMulRepeatM(dst, dst, tempScale, sub_m, qk_n, qk_round_n, 0); - PIPE_BARRIER(V); - } - - __aicore__ __attribute__((always_inline)) inline void QuantPerTokenImpl( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& scale, - uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t pQuantOnline) - { - if (pQuantOnline) { - // scr / scale - TensorDivRepeatM(dst.template ReinterpretCast(), src, scale, sub_m, qk_n, qk_round_n); - } else { - // scr * scale - TensorMulRepeatM(dst.template ReinterpretCast(), src, scale, sub_m, qk_n, qk_round_n, 0); - } - // src fp32 -> casttofp16 -> casttoint8 - uint32_t count = sub_m * qk_round_n; - uint32_t repeat_times = (count + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE; - if (repeat_times < 255) { - conv_v( - dst.template ReinterpretCast(), // dst - dst.template ReinterpretCast(), // src - repeat_times, // repeat_times - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); - } else { - for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { - conv_v( - dst.template ReinterpretCast()[vconv_idx * count / 2], // dst - dst.template ReinterpretCast()[vconv_idx * count / 2], // src - (count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat_times - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); - } - } - PIPE_BARRIER(V); - for (uint32_t row_idx = 0; row_idx < qk_n / HALF_VECTOR_SIZE; ++row_idx) { - AscendC::Cast(dst.template ReinterpretCast()[row_idx * HALF_VECTOR_SIZE], - dst.template ReinterpretCast()[row_idx * HALF_VECTOR_SIZE], AscendC::RoundMode::CAST_RINT, - (uint64_t)0, sub_m, {1, 1, (uint8_t)((qk_round_n) / BLOCK_SIZE), (uint8_t)(qk_round_n / BLOCK_SIZE)}); - } - if (qk_n % HALF_VECTOR_SIZE > 0) { - __set_mask(qk_n % HALF_VECTOR_SIZE); - AscendC::Cast(dst.template ReinterpretCast()[qk_n / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], - dst.template ReinterpretCast()[qk_n / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], AscendC::RoundMode::CAST_RINT, - (uint64_t)0, sub_m, {1, 1, (uint8_t)((qk_round_n) / BLOCK_SIZE), (uint8_t)(qk_round_n / BLOCK_SIZE)}); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); - } - - // SoftmaxStage1 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） - - __aicore__ __attribute__((always_inline)) inline void SoftmaxStage2MLAHeadLoop( - AscendC::GlobalTensor o_tmp_gm_tensor, - AscendC::GlobalTensor go_gm_tensor, - AscendC::GlobalTensor o_gm_tensor, - AscendC::LocalTensor dm32_ubuf_tensor, - AscendC::LocalTensor ll_ubuf_tensor, - AscendC::LocalTensor pm32_ubuf_tensor, - uint32_t n_idx, - uint32_t n_loop, - uint32_t qk_n, - uint32_t qk_round_n, - uint32_t sub_m, - uint64_t o_offset, - uint32_t head_idx, - uint32_t pm_flag_scalar, - uint32_t head_loop, - uint32_t head_loop_idx, - uint32_t q_seq_len, - uint32_t sub_head_num, - uint32_t cur_head_num, - uint32_t numhead_per_process, - uint32_t head_res_row_num, - uint32_t head_start_sblock_idx, - uint32_t tail_res_row_num - ) - { - uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - WAIT_FLAG(V, MTE2, EVENT_ID0); - if (n_idx != 0) { - gm_to_ub( - lo_ubuf_tensor.template ReinterpretCast(), - o_tmp_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - DeQuantPerHeadImpl( - deq_scale_gm_tensor_k1[head_idx], - o_tmp_gm_tensor, - lo_ubuf_tensor, lo_ubuf_tensor.template ReinterpretCast(),// lo_ubuf_tensor use the same ptr - descale_k1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, round_v, round_v, 1, 0); - } - } - SetVectorMask((uint64_t)-1, (uint64_t)-1); - WAIT_FLAG(MTE3, MTE2, EVENT_ID4); - if (n_idx != 0) { - // *** dm = exp(dm) - if (head_loop_idx == 0) { - exp_v(dm32_ubuf_tensor, - dm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - PIPE_BARRIER(V); - // *** gl = dm * gl - mul_v(gl32_ubuf_tensor, - dm32_ubuf_tensor, - gl32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); - // *** gl = ll + gl - add_v(gl32_ubuf_tensor, - gl32_ubuf_tensor, - ll_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); - } - SetVectorMask((uint64_t)-1, (uint64_t)-1); - brcb_v(tv32_ubuf_tensor.ReinterpretCast(), - dm32_ubuf_tensor.ReinterpretCast(), - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - PIPE_BARRIER(V); - if (head_loop > 1) { - gm_to_ub( - go32_ubuf_tensor, - go_gm_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - } - - // *** go = go * dm_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { - mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (__v % FLOAT_VECTOR_SIZE > 0) { - __set_mask(__v % FLOAT_VECTOR_SIZE); - mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); - // *** go = lo + go - add_v(go32_ubuf_tensor, - go32_ubuf_tensor, - lo_ubuf_tensor, - (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); - } else { - // *** gl = ll - if (head_loop_idx == 0) { - ub_to_ub( - gl32_ubuf_tensor, - ll_ubuf_tensor, - 0, // sid - 1, // nBurst - 64 / FLOAT_BLOCK_SIZE, // lenBurst - // round_sub_m / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - PIPE_BARRIER(V); - } - - gm_to_ub( - go32_ubuf_tensor.template ReinterpretCast(), - o_tmp_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - DeQuantPerHeadImpl( - deq_scale_gm_tensor_k1[head_idx], - o_tmp_gm_tensor, - go32_ubuf_tensor, go32_ubuf_tensor.template ReinterpretCast(), - descale_k1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, round_v, round_v, 1, 0); - } else { - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - } - } - SET_FLAG(V, MTE2, EVENT_ID0); - - if (n_idx == n_loop - 1) { - // *** gl_block = expand_to_block(gl) - brcb_v(tv32_ubuf_tensor.ReinterpretCast(), - gl32_ubuf_tensor.ReinterpretCast()[head_loop_idx * 16], - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - PIPE_BARRIER(V); - // *** go = go / gl_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vdiv_idx = 0; vdiv_idx < __v / FLOAT_VECTOR_SIZE; ++vdiv_idx) { - div_v(go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (__v % FLOAT_VECTOR_SIZE > 0) { - __set_mask(__v % FLOAT_VECTOR_SIZE); - div_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); // fix hidden_size=96 - } - PIPE_BARRIER(V); - - // *** go = castfp32to16(go) - conv_v(go_ubuf_tensor, - go32_ubuf_tensor, - (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); - SET_FLAG(V, MTE3, EVENT_ID0); - WAIT_FLAG(V, MTE3, EVENT_ID0); - - uint32_t inner_o_gm_offset = 0; - uint32_t inner_go_ubuf_offset = 0; - - if (head_res_row_num != 0) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset + q_heads * __v * head_start_sblock_idx], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams( - head_res_row_num, // blockCount - __v * 2, // blockLen - 0, // srcStride - __v * (q_heads - 1) * 2, // dstStride - 0 // rsv - ) - ); - inner_o_gm_offset += __v; - inner_go_ubuf_offset += head_res_row_num * __v; - } - - for (uint32_t i = 0; i < numhead_per_process; i++) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams( - q_seq_len, // blockCount - __v * 2, // blockLen - 0, // srcStride - __v * (q_heads - 1) * 2, // dstStride - 0 // rsv - ) - ); - inner_o_gm_offset += __v; - inner_go_ubuf_offset += q_seq_len * __v; - } - - if (tail_res_row_num != 0) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams( - tail_res_row_num, // blockCount - __v * 2, // blockLen - 0, // srcStride - __v * (q_heads - 1) * 2, // dstStride - 0 // rsv - ) - ); - } - // ********************* move O to GM ************************ - if constexpr (IS_RING) { - uint32_t lenBurst = sizeof(OUT_DTYPE); - ln_v(lse32_ubuf_tensor, - gl32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - PIPE_BARRIER(V); - add_v(lse32_ubuf_tensor, - lse32_ubuf_tensor, - gm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); - conv_v(lse_conv_ubuf_tensor, - lse32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); - SET_FLAG(V, MTE3, EVENT_ID1); - WAIT_FLAG(V, MTE3, EVENT_ID1); - // copyout lse - ub_to_gm_align( - lse_gm_tensor[(int64_t)(o_offset / __k)], - lse_conv_ubuf_tensor, - 0, // sid - 1, // nBurst - lenBurst * sub_m * head_loop, // lenBurst - 0, // leftPaddingNum - 0, // rightPaddingNum - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE3, V, EVENT_ID1); - WAIT_FLAG(MTE3, V, EVENT_ID1); - } - - } else if (head_loop > 1) { - SET_FLAG(V, MTE3, EVENT_ID5); - WAIT_FLAG(V, MTE3, EVENT_ID5); - ub_to_gm( - go_gm_tensor, - go32_ubuf_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - } - SET_FLAG(MTE3, MTE2, EVENT_ID4); - } - - __aicore__ __attribute__((always_inline)) inline void SoftmaxStage2MLAHeadLoopTP1( - AscendC::GlobalTensor o_tmp_gm_tensor, - AscendC::GlobalTensor go_gm_tensor, - AscendC::GlobalTensor o_gm_tensor, - AscendC::LocalTensor dm32_ubuf_tensor, - AscendC::LocalTensor ll_ubuf_tensor, - AscendC::LocalTensor pm32_ubuf_tensor, - uint32_t n_idx, - uint32_t n_loop, - uint32_t qk_n, - uint32_t qk_round_n, - uint32_t sub_m, - uint64_t o_offset, - uint32_t head_idx, - uint32_t pm_flag_scalar, - uint32_t head_loop, - uint32_t head_loop_idx, - uint32_t q_seq_len, - uint32_t sub_head_num, - uint32_t cur_head_num, - uint32_t numhead_per_process, - uint32_t head_res_row_num, - uint32_t head_start_sblock_idx, - uint32_t tail_res_row_num - ) - { - uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - WAIT_FLAG(V, MTE2, EVENT_ID0); - if (n_idx != 4) { - gm_to_ub( - lo_ubuf_tensor.template ReinterpretCast(), - o_tmp_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - } - SetVectorMask((uint64_t)-1, (uint64_t)-1); - WAIT_FLAG(MTE3, MTE2, EVENT_ID4); - if (n_idx != 4) { - // expand_to_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - brcb_v(tv32_ubuf_tensor.ReinterpretCast(), - dm32_ubuf_tensor.ReinterpretCast(), - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - PIPE_BARRIER(V); - if (head_loop > 1) { - gm_to_ub( - go32_ubuf_tensor, - go_gm_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - } - - // *** go = go * dm_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { - mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (__v % FLOAT_VECTOR_SIZE > 0) { - __set_mask(__v % FLOAT_VECTOR_SIZE); - mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); - // *** go = lo + go - add_v(go32_ubuf_tensor, - go32_ubuf_tensor, - lo_ubuf_tensor, - (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); - } else { - // *** go = lo - - gm_to_ub( - go32_ubuf_tensor.template ReinterpretCast(), - o_tmp_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - } - SET_FLAG(V, MTE2, EVENT_ID0); - - if (n_idx + 4 > n_loop + 4 - 1) { - // *** gl_block = expand_to_block(gl), 存放于 tv - brcb_v(tv32_ubuf_tensor.ReinterpretCast(), - gl32_ubuf_tensor.ReinterpretCast()[head_loop_idx * 16], - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - PIPE_BARRIER(V); - // *** go = go / gl_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vdiv_idx = 0; vdiv_idx < __v / FLOAT_VECTOR_SIZE; ++vdiv_idx) { - div_v(go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (__v % FLOAT_VECTOR_SIZE > 0) { - __set_mask(__v % FLOAT_VECTOR_SIZE); - div_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); // fix hidden_size=96 - } - PIPE_BARRIER(V); - - // *** go = castfp32to16(go) - conv_v(go_ubuf_tensor, - go32_ubuf_tensor, - (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); - SET_FLAG(V, MTE3, EVENT_ID0); - WAIT_FLAG(V, MTE3, EVENT_ID0); - - uint32_t inner_o_gm_offset = 0; - uint32_t inner_go_ubuf_offset = 0; - - if (head_res_row_num != 0) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset + q_heads * __v * head_start_sblock_idx], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams( - head_res_row_num, // blockCount - __v * 2, // blockLen - 0, // srcStride - __v * (q_heads - 1) * 2, // dstStride - 0 // rsv - ) - ); - inner_o_gm_offset += __v; - inner_go_ubuf_offset += head_res_row_num * __v; - } - - for (uint32_t i = 0; i < numhead_per_process; i++) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams( - q_seq_len, // blockCount - __v * 2, // blockLen - 0, // srcStride - __v * (q_heads - 1) * 2, // dstStride - 0 // rsv - ) - ); - inner_o_gm_offset += __v; - inner_go_ubuf_offset += q_seq_len * __v; - } - - if (tail_res_row_num != 0) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams( - tail_res_row_num, // blockCount - __v * 2, // blockLen - 0, // srcStride - __v * (q_heads - 1) * 2, // dstStride - 0 // rsv - ) - ); - } - // ********************* move O to GM ************************ - if constexpr (IS_RING) { - uint32_t lenBurst = sizeof(OUT_DTYPE); - ln_v(lse32_ubuf_tensor, - gl32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - PIPE_BARRIER(V); - add_v(lse32_ubuf_tensor, - lse32_ubuf_tensor, - gm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); - conv_v(lse_conv_ubuf_tensor, - lse32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); - SET_FLAG(V, MTE3, EVENT_ID1); - WAIT_FLAG(V, MTE3, EVENT_ID1); - // copyout lse - ub_to_gm_align( - lse_gm_tensor[(int64_t)(o_offset / __k)], - lse_conv_ubuf_tensor, - 0, // sid - 1, // nBurst - lenBurst * sub_m * head_loop, // lenBurst - 0, // leftPaddingNum - 0, // rightPaddingNum - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE3, V, EVENT_ID1); - WAIT_FLAG(MTE3, V, EVENT_ID1); - } - } - else if (head_loop > 1) { - SET_FLAG(V, MTE3, EVENT_ID5); - WAIT_FLAG(V, MTE3, EVENT_ID5); - ub_to_gm( - go_gm_tensor, - go32_ubuf_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - } - SET_FLAG(MTE3, MTE2, EVENT_ID4); - } - - __aicore__ __attribute__((always_inline)) inline void TailSoftmaxStage2MLAHeadLoopTP1( - AscendC::GlobalTensor o_tmp_gm_tensor, - AscendC::GlobalTensor go_gm_tensor, - AscendC::GlobalTensor gl_gm_tensor, - AscendC::GlobalTensor gm_gm_tensor, - AscendC::LocalTensor dm32_ubuf_tensor, - AscendC::LocalTensor go32_ubuf_tensor, - AscendC::LocalTensor gl32_ubuf_tensor, - AscendC::LocalTensor gm32_ubuf_tensor, - uint32_t n_idx, - uint32_t n_loop, - uint32_t qk_n, - uint32_t qk_round_n, - uint32_t sub_m, - uint64_t o_offset, - uint32_t head_idx, - uint32_t head_loop, - uint32_t head_loop_idx, - uint32_t q_seq_len, - uint32_t sub_head_num, - uint32_t cur_head_num, - uint32_t numhead_per_process - ) - { - uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - WAIT_FLAG(V, MTE2, EVENT_ID0); - if (n_idx != 4) { - gm_to_ub( - lo_ubuf_tensor.template ReinterpretCast(), - o_tmp_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - } - SetVectorMask((uint64_t)-1, (uint64_t)-1); - WAIT_FLAG(MTE3, MTE2, EVENT_ID4); - if (n_idx != 4) { - // expand_to_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - brcb_v(tv32_ubuf_tensor.ReinterpretCast(), - dm32_ubuf_tensor.ReinterpretCast(), - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - PIPE_BARRIER(V); - if (head_loop > 1) { - gm_to_ub( - go32_ubuf_tensor, - go_gm_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - } - - // *** go = go * dm_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { - mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (__v % FLOAT_VECTOR_SIZE > 0) { - __set_mask(__v % FLOAT_VECTOR_SIZE); - mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); - // *** go = lo + go - add_v(go32_ubuf_tensor, - go32_ubuf_tensor, - lo_ubuf_tensor, - (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); - } else { - // *** go = lo - - gm_to_ub( - go32_ubuf_tensor.template ReinterpretCast(), - o_tmp_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - } - SET_FLAG(V, MTE2, EVENT_ID0); + // TensorSubValueRepeatM / TensorDivRepeatM / TensorMulRepeatM / DeQuantPerHeadImpl / QuantPerTokenImpl + // 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） - if (n_idx + 4 > n_loop + 4 - 1) { - // The last step to process the o with dividing and copyout - // TODO: Maybe the following two don't need waiting - // Copyout gl32_ubuf_tensor to gl_gm_tensor - ub_to_gm( - gl_gm_tensor, - gl32_ubuf_tensor, - 0, - 1, - sub_m / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - - // Copyout rowmax to global gm32_ubuf_tensor - ub_to_gm( - gm_gm_tensor, - gm32_ubuf_tensor, - 0, - 1, - sub_m / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - // Copyout go32_ubuf_tensor to go_gm_tensor - // This is needed to wait for former calculation - SET_FLAG(V, MTE3, EVENT_ID5); - WAIT_FLAG(V, MTE3, EVENT_ID5); - ub_to_gm( - go_gm_tensor, - go32_ubuf_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - - } - else if (head_loop > 1) { - SET_FLAG(V, MTE3, EVENT_ID5); - WAIT_FLAG(V, MTE3, EVENT_ID5); - ub_to_gm( - go_gm_tensor, - go32_ubuf_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - } - SET_FLAG(MTE3, MTE2, EVENT_ID4); - PIPE_BARRIER(ALL); - } + // SoftmaxStage2MLAHeadLoop / SoftmaxStage2MLAHeadLoopTP1 / TailSoftmaxStage2MLAHeadLoopTP1 + // 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） __aicore__ __attribute((always_inline)) inline void SoftmaxGatherTP1( AscendC::GlobalTensor o_gm_tensor, diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h index 85d04ce..f307f14 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h @@ -10,6 +10,176 @@ // 本文件仅被 MLADecoderAiv 类 include，包含 Vector 侧业务调度函数。 // AIC 侧业务函数在 multi_latent_attention_bs.h 中。 +// ====== AIV refactor: context struct + init functions ====== + +// 非 TP1 路径上下文 +struct VectorContext { + // input params + uint32_t cur_batch; + uint32_t start_head; + uint32_t cur_nIndx; + uint32_t cur_q_seqlen; + uint32_t cur_kv_seqlen; + uint32_t cur_head_num; + uint32_t offset_tiling; + + // addresses + uint64_t addr_o_scalar; + uint64_t addr_mask_scalar; + uint32_t mask_offset; + + // loop & size + uint32_t pp_n_scalar; + uint32_t sub_n_loop; + uint32_t real_n_loop; + uint32_t n_loop; + + // QK dims + uint32_t qk_n; + uint32_t qk_round_n; + uint32_t qk_n_2; + uint32_t qk_round_n_2; + + // head split + uint32_t sub_head_num; + uint32_t sub_m; + uint32_t head_idx; + uint64_t o_offset; + + // tail info + uint32_t tail_len; + bool prev_tail_mask; +}; + +__aicore__ __attribute__((always_inline)) inline void InitVectorContext( + VectorContext &ctx, uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, + uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, + uint32_t offset_tiling) +{ + ctx.cur_batch = cur_batch; + ctx.start_head = start_head; + ctx.cur_nIndx = cur_nIndx; + ctx.cur_q_seqlen = cur_q_seqlen; + ctx.cur_kv_seqlen = cur_kv_seqlen; + ctx.cur_head_num = cur_head_num; + ctx.offset_tiling = offset_tiling; + + // addr + uint32_t addr_o_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 4 + offset_tiling)); + uint32_t addr_o_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 5 + offset_tiling)); + ctx.addr_o_scalar = (uint64_t)(((uint64_t)addr_o_high32) << 32 | addr_o_loww32); + + uint32_t addr_mask_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 6 + offset_tiling)); + uint32_t addr_mask_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 7 + offset_tiling)); + ctx.addr_mask_scalar = (uint64_t)(((uint64_t)addr_mask_high32) << 32 | addr_mask_loww32); + ctx.mask_offset = ctx.addr_mask_scalar; + + // loop & size + ctx.pp_n_scalar = block_size; + ctx.sub_n_loop = ctx.pp_n_scalar / block_size; + ctx.real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; + ctx.n_loop = (cur_kv_seqlen + ctx.pp_n_scalar - 1) / ctx.pp_n_scalar; + + // QK dims + ctx.qk_n = ctx.pp_n_scalar; + ctx.qk_round_n = RoundUp(ctx.qk_n); + ctx.qk_n_2 = ctx.pp_n_scalar; + ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); + + // head split + ctx.sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; + ctx.sub_m = ctx.sub_head_num * cur_q_seqlen; + ctx.head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; + ctx.o_offset = ctx.addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; + + // tail info + ctx.tail_len = cur_kv_seqlen - (ctx.n_loop - 1) * ctx.pp_n_scalar; + ctx.prev_tail_mask = (ctx.n_loop > 1 && ctx.tail_len < cur_q_seqlen - 1); +} + +// TP1 路径上下文 +struct VectorTP1Context { + // input params + uint32_t cur_batch; + uint32_t start_head; + uint32_t cur_nIndx; + uint32_t cur_q_seqlen; + uint32_t cur_kv_seqlen; + uint32_t cur_head_num; + uint32_t offset_tiling; + + // addresses + uint64_t addr_o_scalar; + uint64_t addr_mask_scalar; + uint32_t mask_offset; + + // loop & size + uint32_t pp_n_scalar; + uint32_t sub_n_loop; + uint32_t real_n_loop; + uint32_t n_loop; + + // QK dims + uint32_t qk_n; + uint32_t qk_round_n; + uint32_t qk_n_2; + uint32_t qk_round_n_2; + + // head split + uint32_t sub_head_num; + uint32_t sub_m; + uint32_t head_idx; + uint64_t o_offset; + + // TP1 specific + uint32_t s_block_stack; + uint32_t m_slice; + uint32_t m_end; +}; + +__aicore__ __attribute__((always_inline)) inline void InitVectorTP1Context( + VectorTP1Context &ctx, uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, + uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, + uint32_t offset_tiling) +{ + ctx.cur_batch = cur_batch; + ctx.start_head = start_head; + ctx.cur_nIndx = cur_nIndx; + ctx.cur_q_seqlen = cur_q_seqlen; + ctx.cur_kv_seqlen = cur_kv_seqlen; + ctx.cur_head_num = cur_head_num; + ctx.offset_tiling = offset_tiling; + + // addr + uint32_t prev_task = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); + ctx.addr_o_scalar = prev_task * q_heads * embedding_size; + ctx.addr_mask_scalar = 0; + ctx.mask_offset = 0; + + // loop & size + ctx.pp_n_scalar = block_size; + ctx.sub_n_loop = ctx.pp_n_scalar / block_size; + ctx.real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; + ctx.n_loop = (cur_kv_seqlen + ctx.pp_n_scalar - 1) / ctx.pp_n_scalar; + + // QK dims + ctx.qk_n = ctx.pp_n_scalar; + ctx.qk_round_n = RoundUp(ctx.qk_n); + ctx.qk_n_2 = ctx.pp_n_scalar; + ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); + + // head split + ctx.sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; + ctx.sub_m = ctx.sub_head_num * cur_q_seqlen; + ctx.head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; + ctx.o_offset = ctx.addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; + + // TP1 specific + ctx.s_block_stack = 4; + ctx.m_slice = FLOAT_VECTOR_SIZE / ctx.s_block_stack; + ctx.m_end = (ctx.sub_m + ctx.m_slice - 1) / ctx.m_slice; +} + // ====== 非 TP1 路径：Phase 1 / Phase 2 业务调度 ====== // 业务子函数1a：QK 数据加载（INT8 路径） @@ -461,6 +631,370 @@ __aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage1( PlatformSoftmaxStage1PostSync(); } +// 业务函数：非TP1 Phase 2 — SoftmaxStage2MLAHeadLoop 核心 Head Loop 计算 +// 包含 load lo → DeQuant → exp(dm) → gl=dm*gl → gl=ll+gl → brcb dm → go=go*dm → go=lo+go +// → brcb gl → go=go/gl → conv_v → DataCopyPad输出 → Ring LSE copyout +__aicore__ __attribute__((always_inline)) inline void SoftmaxStage2MLAHeadLoop( + AscendC::GlobalTensor o_tmp_gm_tensor, + AscendC::GlobalTensor go_gm_tensor, + AscendC::GlobalTensor o_gm_tensor, + AscendC::LocalTensor dm32_ubuf_tensor, + AscendC::LocalTensor ll_ubuf_tensor, + AscendC::LocalTensor pm32_ubuf_tensor, + uint32_t n_idx, + uint32_t n_loop, + uint32_t qk_n, + uint32_t qk_round_n, + uint32_t sub_m, + uint64_t o_offset, + uint32_t head_idx, + uint32_t pm_flag_scalar, + uint32_t head_loop, + uint32_t head_loop_idx, + uint32_t q_seq_len, + uint32_t sub_head_num, + uint32_t cur_head_num, + uint32_t numhead_per_process, + uint32_t head_res_row_num, + uint32_t head_start_sblock_idx, + uint32_t tail_res_row_num + ) +{ + uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; + WAIT_FLAG(V, MTE2, EVENT_ID0); + if (n_idx != 0) { + gm_to_ub( + lo_ubuf_tensor.template ReinterpretCast(), + o_tmp_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + DeQuantPerHeadImpl( + deq_scale_gm_tensor_k1[head_idx], + o_tmp_gm_tensor, + lo_ubuf_tensor, lo_ubuf_tensor.template ReinterpretCast(),// lo_ubuf_tensor use the same ptr + descale_k1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, round_v, round_v, 1, 0); + } + } + SetVectorMask((uint64_t)-1, (uint64_t)-1); + WAIT_FLAG(MTE3, MTE2, EVENT_ID4); + if (n_idx != 0) { + // *** dm = exp(dm) + if (head_loop_idx == 0) { + exp_v(dm32_ubuf_tensor, + dm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + // *** gl = dm * gl + mul_v(gl32_ubuf_tensor, + dm32_ubuf_tensor, + gl32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + // *** gl = ll + gl + add_v(gl32_ubuf_tensor, + gl32_ubuf_tensor, + ll_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + } + SetVectorMask((uint64_t)-1, (uint64_t)-1); + brcb_v(tv32_ubuf_tensor.ReinterpretCast(), + dm32_ubuf_tensor.ReinterpretCast(), + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + if (head_loop > 1) { + gm_to_ub( + go32_ubuf_tensor, + go_gm_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + + // *** go = go * dm_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { + mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (__v % FLOAT_VECTOR_SIZE > 0) { + __set_mask(__v % FLOAT_VECTOR_SIZE); + mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + // *** go = lo + go + add_v(go32_ubuf_tensor, + go32_ubuf_tensor, + lo_ubuf_tensor, + (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + } else { + // *** gl = ll + if (head_loop_idx == 0) { + ub_to_ub( + gl32_ubuf_tensor, + ll_ubuf_tensor, + 0, // sid + 1, // nBurst + 64 / FLOAT_BLOCK_SIZE, // lenBurst + // round_sub_m / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + PIPE_BARRIER(V); + } + + gm_to_ub( + go32_ubuf_tensor.template ReinterpretCast(), + o_tmp_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + DeQuantPerHeadImpl( + deq_scale_gm_tensor_k1[head_idx], + o_tmp_gm_tensor, + go32_ubuf_tensor, go32_ubuf_tensor.template ReinterpretCast(), + descale_k1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, round_v, round_v, 1, 0); + } else { + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + } + SET_FLAG(V, MTE2, EVENT_ID0); + + if (n_idx == n_loop - 1) { + // *** gl_block = expand_to_block(gl) + brcb_v(tv32_ubuf_tensor.ReinterpretCast(), + gl32_ubuf_tensor.ReinterpretCast()[head_loop_idx * 16], + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + // *** go = go / gl_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vdiv_idx = 0; vdiv_idx < __v / FLOAT_VECTOR_SIZE; ++vdiv_idx) { + div_v(go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (__v % FLOAT_VECTOR_SIZE > 0) { + __set_mask(__v % FLOAT_VECTOR_SIZE); + div_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); // fix hidden_size=96 + } + PIPE_BARRIER(V); + + // *** go = castfp32to16(go) + conv_v(go_ubuf_tensor, + go32_ubuf_tensor, + (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + SET_FLAG(V, MTE3, EVENT_ID0); + WAIT_FLAG(V, MTE3, EVENT_ID0); + + uint32_t inner_o_gm_offset = 0; + uint32_t inner_go_ubuf_offset = 0; + + if (head_res_row_num != 0) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset + q_heads * __v * head_start_sblock_idx], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams( + head_res_row_num, // blockCount + __v * 2, // blockLen + 0, // srcStride + __v * (q_heads - 1) * 2, // dstStride + 0 // rsv + ) + ); + inner_o_gm_offset += __v; + inner_go_ubuf_offset += head_res_row_num * __v; + } + + for (uint32_t i = 0; i < numhead_per_process; i++) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams( + q_seq_len, // blockCount + __v * 2, // blockLen + 0, // srcStride + __v * (q_heads - 1) * 2, // dstStride + 0 // rsv + ) + ); + inner_o_gm_offset += __v; + inner_go_ubuf_offset += q_seq_len * __v; + } + + if (tail_res_row_num != 0) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams( + tail_res_row_num, // blockCount + __v * 2, // blockLen + 0, // srcStride + __v * (q_heads - 1) * 2, // dstStride + 0 // rsv + ) + ); + } + // ********************* move O to GM ************************ + if constexpr (IS_RING) { + uint32_t lenBurst = sizeof(OUT_DTYPE); + ln_v(lse32_ubuf_tensor, + gl32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + add_v(lse32_ubuf_tensor, + lse32_ubuf_tensor, + gm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + conv_v(lse_conv_ubuf_tensor, + lse32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + SET_FLAG(V, MTE3, EVENT_ID1); + WAIT_FLAG(V, MTE3, EVENT_ID1); + // copyout lse + ub_to_gm_align( + lse_gm_tensor[(int64_t)(o_offset / __k)], + lse_conv_ubuf_tensor, + 0, // sid + 1, // nBurst + lenBurst * sub_m * head_loop, // lenBurst + 0, // leftPaddingNum + 0, // rightPaddingNum + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE3, V, EVENT_ID1); + WAIT_FLAG(MTE3, V, EVENT_ID1); + } + + } else if (head_loop > 1) { + SET_FLAG(V, MTE3, EVENT_ID5); + WAIT_FLAG(V, MTE3, EVENT_ID5); + ub_to_gm( + go_gm_tensor, + go32_ubuf_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + } + SET_FLAG(MTE3, MTE2, EVENT_ID4); +} + // 业务函数：非TP1 Phase 2 — Softmax Stage2 调度 // 包含 head 循环计算、WaitFlag、SoftmaxStage2MLAHeadLoop 调用 __aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage2( @@ -638,6 +1172,304 @@ __aicore__ __attribute__((always_inline)) inline void ScheduleOnlineSoftmaxStage FftsCrossCoreSync(SOFTMAX_READY_DECODER); } +// 业务函数：TP1 Phase 2 — SoftmaxStage2MLAHeadLoopTP1 核心 Head Loop 计算 +// 与非TP1相似，差异：n_idx!=4 vs n_idx!=0，无 DeQuant/exp(dm)/gl=dm*gl/gl=ll+gl 路径 +__aicore__ __attribute__((always_inline)) inline void SoftmaxStage2MLAHeadLoopTP1( + AscendC::GlobalTensor o_tmp_gm_tensor, + AscendC::GlobalTensor go_gm_tensor, + AscendC::GlobalTensor o_gm_tensor, + AscendC::LocalTensor dm32_ubuf_tensor, + AscendC::LocalTensor ll_ubuf_tensor, + AscendC::LocalTensor pm32_ubuf_tensor, + uint32_t n_idx, + uint32_t n_loop, + uint32_t qk_n, + uint32_t qk_round_n, + uint32_t sub_m, + uint64_t o_offset, + uint32_t head_idx, + uint32_t pm_flag_scalar, + uint32_t head_loop, + uint32_t head_loop_idx, + uint32_t q_seq_len, + uint32_t sub_head_num, + uint32_t cur_head_num, + uint32_t numhead_per_process, + uint32_t head_res_row_num, + uint32_t head_start_sblock_idx, + uint32_t tail_res_row_num + ) +{ + uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; + WAIT_FLAG(V, MTE2, EVENT_ID0); + if (n_idx != 4) { + gm_to_ub( + lo_ubuf_tensor.template ReinterpretCast(), + o_tmp_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + SetVectorMask((uint64_t)-1, (uint64_t)-1); + WAIT_FLAG(MTE3, MTE2, EVENT_ID4); + if (n_idx != 4) { + // expand_to_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + brcb_v(tv32_ubuf_tensor.ReinterpretCast(), + dm32_ubuf_tensor.ReinterpretCast(), + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + if (head_loop > 1) { + gm_to_ub( + go32_ubuf_tensor, + go_gm_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + + // *** go = go * dm_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { + mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (__v % FLOAT_VECTOR_SIZE > 0) { + __set_mask(__v % FLOAT_VECTOR_SIZE); + mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + // *** go = lo + go + add_v(go32_ubuf_tensor, + go32_ubuf_tensor, + lo_ubuf_tensor, + (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + } else { + // *** go = lo + + gm_to_ub( + go32_ubuf_tensor.template ReinterpretCast(), + o_tmp_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + SET_FLAG(V, MTE2, EVENT_ID0); + + if (n_idx + 4 > n_loop + 4 - 1) { + // *** gl_block = expand_to_block(gl), 存放于 tv + brcb_v(tv32_ubuf_tensor.ReinterpretCast(), + gl32_ubuf_tensor.ReinterpretCast()[head_loop_idx * 16], + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + // *** go = go / gl_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vdiv_idx = 0; vdiv_idx < __v / FLOAT_VECTOR_SIZE; ++vdiv_idx) { + div_v(go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (__v % FLOAT_VECTOR_SIZE > 0) { + __set_mask(__v % FLOAT_VECTOR_SIZE); + div_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); // fix hidden_size=96 + } + PIPE_BARRIER(V); + + // *** go = castfp32to16(go) + conv_v(go_ubuf_tensor, + go32_ubuf_tensor, + (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + SET_FLAG(V, MTE3, EVENT_ID0); + WAIT_FLAG(V, MTE3, EVENT_ID0); + + uint32_t inner_o_gm_offset = 0; + uint32_t inner_go_ubuf_offset = 0; + + if (head_res_row_num != 0) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset + q_heads * __v * head_start_sblock_idx], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams( + head_res_row_num, // blockCount + __v * 2, // blockLen + 0, // srcStride + __v * (q_heads - 1) * 2, // dstStride + 0 // rsv + ) + ); + inner_o_gm_offset += __v; + inner_go_ubuf_offset += head_res_row_num * __v; + } + + for (uint32_t i = 0; i < numhead_per_process; i++) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams( + q_seq_len, // blockCount + __v * 2, // blockLen + 0, // srcStride + __v * (q_heads - 1) * 2, // dstStride + 0 // rsv + ) + ); + inner_o_gm_offset += __v; + inner_go_ubuf_offset += q_seq_len * __v; + } + + if (tail_res_row_num != 0) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams( + tail_res_row_num, // blockCount + __v * 2, // blockLen + 0, // srcStride + __v * (q_heads - 1) * 2, // dstStride + 0 // rsv + ) + ); + } + // ********************* move O to GM ************************ + if constexpr (IS_RING) { + uint32_t lenBurst = sizeof(OUT_DTYPE); + ln_v(lse32_ubuf_tensor, + gl32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + add_v(lse32_ubuf_tensor, + lse32_ubuf_tensor, + gm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + conv_v(lse_conv_ubuf_tensor, + lse32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + SET_FLAG(V, MTE3, EVENT_ID1); + WAIT_FLAG(V, MTE3, EVENT_ID1); + // copyout lse + ub_to_gm_align( + lse_gm_tensor[(int64_t)(o_offset / __k)], + lse_conv_ubuf_tensor, + 0, // sid + 1, // nBurst + lenBurst * sub_m * head_loop, // lenBurst + 0, // leftPaddingNum + 0, // rightPaddingNum + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE3, V, EVENT_ID1); + WAIT_FLAG(MTE3, V, EVENT_ID1); + } + } + else if (head_loop > 1) { + SET_FLAG(V, MTE3, EVENT_ID5); + WAIT_FLAG(V, MTE3, EVENT_ID5); + ub_to_gm( + go_gm_tensor, + go32_ubuf_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + } + SET_FLAG(MTE3, MTE2, EVENT_ID4); +} + // 业务函数：TP1 Phase 2 — SoftmaxStage2MLAHeadLoopTP1 调度（非 Tail 路径） // 包含 head 循环计算、WaitFlag、SoftmaxStage2MLAHeadLoopTP1 调用 __aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage2TP1( @@ -697,6 +1529,191 @@ __aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage2TP1( } } +// 业务函数：TP1 Phase 2 (Tail) — TailSoftmaxStage2MLAHeadLoopTP1 核心 Head Loop 计算 +// 与 TP1 版本相似，但增加额外的 gl/gm/go copyout 逻辑 +__aicore__ __attribute__((always_inline)) inline void TailSoftmaxStage2MLAHeadLoopTP1( + AscendC::GlobalTensor o_tmp_gm_tensor, + AscendC::GlobalTensor go_gm_tensor, + AscendC::GlobalTensor gl_gm_tensor, + AscendC::GlobalTensor gm_gm_tensor, + AscendC::LocalTensor dm32_ubuf_tensor, + AscendC::LocalTensor go32_ubuf_tensor, + AscendC::LocalTensor gl32_ubuf_tensor, + AscendC::LocalTensor gm32_ubuf_tensor, + uint32_t n_idx, + uint32_t n_loop, + uint32_t qk_n, + uint32_t qk_round_n, + uint32_t sub_m, + uint64_t o_offset, + uint32_t head_idx, + uint32_t head_loop, + uint32_t head_loop_idx, + uint32_t q_seq_len, + uint32_t sub_head_num, + uint32_t cur_head_num, + uint32_t numhead_per_process + ) +{ + uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; + WAIT_FLAG(V, MTE2, EVENT_ID0); + if (n_idx != 4) { + gm_to_ub( + lo_ubuf_tensor.template ReinterpretCast(), + o_tmp_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + SetVectorMask((uint64_t)-1, (uint64_t)-1); + WAIT_FLAG(MTE3, MTE2, EVENT_ID4); + if (n_idx != 4) { + // expand_to_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + brcb_v(tv32_ubuf_tensor.ReinterpretCast(), + dm32_ubuf_tensor.ReinterpretCast(), + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + if (head_loop > 1) { + gm_to_ub( + go32_ubuf_tensor, + go_gm_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + + // *** go = go * dm_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { + mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (__v % FLOAT_VECTOR_SIZE > 0) { + __set_mask(__v % FLOAT_VECTOR_SIZE); + mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + // *** go = lo + go + add_v(go32_ubuf_tensor, + go32_ubuf_tensor, + lo_ubuf_tensor, + (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + } else { + // *** go = lo + + gm_to_ub( + go32_ubuf_tensor.template ReinterpretCast(), + o_tmp_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + SET_FLAG(V, MTE2, EVENT_ID0); + + if (n_idx + 4 > n_loop + 4 - 1) { + // The last step to process the o with dividing and copyout + // TODO: Maybe the following two don't need waiting + // Copyout gl32_ubuf_tensor to gl_gm_tensor + ub_to_gm( + gl_gm_tensor, + gl32_ubuf_tensor, + 0, + 1, + sub_m / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + + // Copyout rowmax to global gm32_ubuf_tensor + ub_to_gm( + gm_gm_tensor, + gm32_ubuf_tensor, + 0, + 1, + sub_m / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + // Copyout go32_ubuf_tensor to go_gm_tensor + // This is needed to wait for former calculation + SET_FLAG(V, MTE3, EVENT_ID5); + WAIT_FLAG(V, MTE3, EVENT_ID5); + ub_to_gm( + go_gm_tensor, + go32_ubuf_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + + } + else if (head_loop > 1) { + SET_FLAG(V, MTE3, EVENT_ID5); + WAIT_FLAG(V, MTE3, EVENT_ID5); + ub_to_gm( + go_gm_tensor, + go32_ubuf_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + } + SET_FLAG(MTE3, MTE2, EVENT_ID4); + PIPE_BARRIER(ALL); +} + // 业务函数：TP1 Phase 2 (Tail) — TailSoftmaxStage2MLAHeadLoopTP1 调度（Tail 路径） // 包含 head 循环计算、WaitFlag、TailSoftmaxStage2MLAHeadLoopTP1 调用 __aicore__ __attribute__((always_inline)) inline void ScheduleTailSoftmaxStage2TP1( @@ -807,6 +1824,56 @@ __aicore__ __attribute__((always_inline)) inline void TailInnerRunVectorChangeTP // ====== Tensor 逐行重复计算（Vector 业务基础函数）====== +// Tensor 逐行减法（brcb 广播 Max + sub_v 循环 + tail 处理） +__aicore__ __attribute__((always_inline)) inline void TensorSubValueRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& MaxTensor, + const AscendC::LocalTensor& tempMaxTensor, + uint32_t sub_m, + uint32_t round_sub_m, + uint32_t qk_n, + uint32_t qk_round_n) +{ + brcb_v( + tempMaxTensor.ReinterpretCast(), + MaxTensor.ReinterpretCast(), + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + for (uint32_t sub_v_idx = 0; sub_v_idx < qk_n / FLOAT_VECTOR_SIZE; ++sub_v_idx) { + sub_v(dst[sub_v_idx * FLOAT_VECTOR_SIZE], + src[sub_v_idx * FLOAT_VECTOR_SIZE], + tempMaxTensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + sub_v(dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tempMaxTensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); +} + // Tensor 逐行除法（div_v 循环 + tail 处理） __aicore__ __attribute__((always_inline)) inline void TensorDivRepeatM( const AscendC::LocalTensor& dst, @@ -846,6 +1913,144 @@ __aicore__ __attribute__((always_inline)) inline void TensorDivRepeatM( PIPE_BARRIER(V); } +// Tensor 行归约求最大值（cmax_v + max_v 循环 + tail 处理） +__aicore__ __attribute__((always_inline)) inline void ReduceMaxRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& tempTensor, + uint32_t sub_m, + uint32_t qk_n, + uint32_t qk_round_n) +{ + if (qk_n <= FLOAT_VECTOR_SIZE) { + __set_mask(qk_n); + cmax_v(dst, + src, + sub_m, // repeat + 1, // dstRepeatStride + 1, // srcBlockStride + qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride + ); + } else { + ub_to_ub( + tempTensor, + src, + 0, // sid + sub_m, // nBurst + HALF_VECTOR_SIZE / BLOCK_SIZE, // lenBurst + (qk_round_n - FLOAT_VECTOR_SIZE) / FLOAT_BLOCK_SIZE, // srcGap + 0 // dstGap + ); + PIPE_BARRIER(V); + for (uint32_t rowmax_idx = 1; rowmax_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowmax_idx) { + max_v( + tempTensor, + tempTensor, + src[rowmax_idx * FLOAT_VECTOR_SIZE], + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride + ); + PIPE_BARRIER(V); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + max_v( + tempTensor, + tempTensor, + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride + ); + } + PIPE_BARRIER(V); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + cmax_v( + dst, + tempTensor, + sub_m, // repeat + 1, // dstRepeatStride + 1, // srcBlockStride + 8 // srcRepeatStride + ); + } + SetVectorMask((uint64_t)-1, (uint64_t)-1); + PIPE_BARRIER(V); +} + +// Tensor 行归约求和（cadd_v + add_v 循环 + tail 处理） +__aicore__ __attribute__((always_inline)) inline void ReduceSumRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + uint32_t sub_m, + uint32_t qk_n, + uint32_t qk_round_n) +{ + if (qk_n <= FLOAT_VECTOR_SIZE) { + __set_mask(qk_n); + cadd_v( + dst, + src, + sub_m, // repeat + 1, // dstRepeatStride + 1, // srcBlockStride + qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else { + for (uint32_t rowsum_idx = 1; rowsum_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowsum_idx) { + add_v( + src, + src, + src[rowsum_idx * FLOAT_VECTOR_SIZE], + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride + ); + PIPE_BARRIER(V); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + add_v( + src, + src, + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + + cadd_v( + dst, + src, + sub_m, // repeat + 1, // dstRepeatStride + 1, // srcBlockStride + qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride + ); + } +} + // Tensor 逐行乘法（mul_v 循环 + tail 处理） __aicore__ __attribute__((always_inline)) inline void TensorMulRepeatM( const AscendC::LocalTensor& dst, From 8915d1d74e5a8bd7ad27dad86a8da61693d5c5d5 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Thu, 13 Aug 2026 14:09:46 +0800 Subject: [PATCH 17/24] feat: you hua mla --- .../multi_latent_attention_aiv_arch32.h | 970 ++++++++++- .../op_kernel/multi_latent_attention_aiv_bs.h | 1454 +++-------------- 2 files changed, 1207 insertions(+), 1217 deletions(-) diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h index 54bc73c..d1fbf15 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h @@ -116,18 +116,972 @@ __aicore__ __attribute__((always_inline)) inline void PlatformWaitVectorPipeSync WAIT_FLAG(V, MTE2, EVENT_ID2); } -// 平台函数：MTE2→V EVENT_ID0 同步（SET+WAIT） -// 用于 DeQuantPerHeadImpl：src 搬入后通知 V 管道可读 -__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSyncEvent0() +// 平台函数：向量广播乘法 go = go * dm_block +// 封装 brcb dm → mul 循环 → mask 尾部 → PIPE_BARRIER +// 用于 Stage2 段落B：go = go * dm_block（三函数完全一致） +__aicore__ __attribute__((always_inline)) inline void PlatformMulVectorByBroadcast( + AscendC::LocalTensor go32_ubuf_tensor, + AscendC::LocalTensor tv32_ubuf_tensor, + AscendC::LocalTensor dm32_ubuf_tensor, + uint32_t sub_m, + uint32_t round_sub_m, + uint32_t round_v, + uint32_t __v) +{ + // expand_to_block: dm → tv + SetVectorMask((uint64_t)-1, (uint64_t)-1); + brcb_v(tv32_ubuf_tensor, + dm32_ubuf_tensor.ReinterpretCast(), + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + + // go = go * dm_block (full vector iterations) + SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { + mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + // tail mask iteration + if (__v % FLOAT_VECTOR_SIZE > 0) { + __set_mask(__v % FLOAT_VECTOR_SIZE); + mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); +} + +// 平台函数：向量广播除法 go = go / gl_block +// 封装 brcb gl → div 循环 → mask 尾部 → PIPE_BARRIER +// 用于 Stage2 段落C：go = go / gl_block（非TP1与TP1完全一致） +__aicore__ __attribute__((always_inline)) inline void PlatformDivVectorByBroadcast( + AscendC::LocalTensor go32_ubuf_tensor, + AscendC::LocalTensor tv32_ubuf_tensor, + AscendC::LocalTensor gl32_ubuf_tensor, + uint32_t sub_m, + uint32_t round_sub_m, + uint32_t round_v, + uint32_t __v, + uint32_t head_loop_idx) +{ + // gl_block = expand_to_block(gl), stored in tv + brcb_v(tv32_ubuf_tensor, + gl32_ubuf_tensor.ReinterpretCast()[head_loop_idx * 16], + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + + // go = go / gl_block (full vector iterations) + SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vdiv_idx = 0; vdiv_idx < __v / FLOAT_VECTOR_SIZE; ++vdiv_idx) { + div_v(go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + // tail mask iteration + if (__v % FLOAT_VECTOR_SIZE > 0) { + __set_mask(__v % FLOAT_VECTOR_SIZE); + div_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); // fix hidden_size=96 + } + PIPE_BARRIER(V); +} + +// 平台函数：中间结果保存（段落D） +// 封装 SET_FLAG(V,MTE3,E5) → WAIT_FLAG → ub_to_gm go32 +// 用于 Stage2 head_loop>1 且非最后一轮时保存 go32 到 GM +// 三函数（非TP1/TP1/Tail）完全一致 +__aicore__ __attribute__((always_inline)) inline void PlatformIntermediateSave( + AscendC::GlobalTensor go_gm_tensor, + AscendC::LocalTensor go32_ubuf_tensor, + uint32_t sub_m, + uint32_t round_v) { + SET_FLAG(V, MTE3, EVENT_ID5); + WAIT_FLAG(V, MTE3, EVENT_ID5); + ub_to_gm( + go_gm_tensor, + go32_ubuf_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); +} + +// 平台函数：从 GM 搬入 go 到 UB（段落B 使用） +// 封装 gm_to_ub + SET_FLAG/WAIT_FLAG(MTE2, V, EVENT_ID0) +// 用于 Stage2MergeAccumulate head_loop>1 时加载上一轮 go +__aicore__ __attribute__((always_inline)) inline void PlatformLoadGoFromGM( + AscendC::GlobalTensor go_gm_tensor, + AscendC::LocalTensor go32_ubuf_tensor, + uint32_t sub_m, + uint32_t round_v) +{ + gm_to_ub( + go32_ubuf_tensor, go_gm_tensor, + 0, 1, sub_m * round_v / FLOAT_BLOCK_SIZE, 0, 0); SET_FLAG(MTE2, V, EVENT_ID0); WAIT_FLAG(MTE2, V, EVENT_ID0); } -// 平台函数：MTE2→V EVENT_ID2 同步（SET+WAIT） -// 用于 DeQuantPerHeadImpl：deScale 搬入后通知 V 管道可做 online 乘法 -__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSyncEvent2() +// 平台函数：向量加法 go = lo + go（段落B 使用） +// 封装 add_v + PIPE_BARRIER(V) +// 用于 Stage2MergeAccumulate 累加 lo 到 go +__aicore__ __attribute__((always_inline)) inline void PlatformAddLoToGo( + AscendC::LocalTensor go32_ubuf_tensor, + AscendC::LocalTensor lo_ubuf_tensor, + uint32_t sub_m, + uint32_t round_v) { - SET_FLAG(MTE2, V, EVENT_ID2); - WAIT_FLAG(MTE2, V, EVENT_ID2); + add_v(go32_ubuf_tensor, + go32_ubuf_tensor, lo_ubuf_tensor, + (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, 1, 1, 8, 8, 8); + PIPE_BARRIER(V); +} + +// 平台函数：FP32→FP16/BF16 类型转换 + V→MTE3 同步（段落C 使用） +// 封装 conv_v + SET_FLAG/WAIT_FLAG(V, MTE3, EVENT_ID0) +// 用于 Stage2FinalizeAndOutput go=castfp32to16(go) +__aicore__ __attribute__((always_inline)) inline void PlatformConvGoToOutput( + AscendC::LocalTensor go_ubuf_tensor, + AscendC::LocalTensor go32_ubuf_tensor, + uint32_t sub_m, + uint32_t round_v) +{ + conv_v(go_ubuf_tensor, + go32_ubuf_tensor, + (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, 1, 4, 8); + SET_FLAG(V, MTE3, EVENT_ID0); + WAIT_FLAG(V, MTE3, EVENT_ID0); +} + +// 平台函数：DataCopyPad 输出到 GM（段落C 使用） +// 封装 head_res_row_num / numhead_per_process / tail_res_row_num 三段 DataCopyPad +// 用于 Stage2FinalizeAndOutput 输出 attention 结果到 o_gm +__aicore__ __attribute__((always_inline)) inline void PlatformDataCopyPadOutput( + AscendC::GlobalTensor o_gm_tensor, + AscendC::LocalTensor go_ubuf_tensor, + uint32_t sub_m, + uint32_t round_v, + uint32_t q_seq_len, + uint32_t numhead_per_process, + uint32_t head_res_row_num, + uint32_t head_start_sblock_idx, + uint32_t tail_res_row_num) +{ + uint32_t inner_o_gm_offset = 0; + uint32_t inner_go_ubuf_offset = 0; + + if (head_res_row_num != 0) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset + q_heads * __v * head_start_sblock_idx], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams(head_res_row_num, __v * 2, 0, + __v * (q_heads - 1) * 2, 0)); + inner_o_gm_offset += __v; + inner_go_ubuf_offset += head_res_row_num * __v; + } + + for (uint32_t i = 0; i < numhead_per_process; i++) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams(q_seq_len, __v * 2, 0, + __v * (q_heads - 1) * 2, 0)); + inner_o_gm_offset += __v; + inner_go_ubuf_offset += q_seq_len * __v; + } + + if (tail_res_row_num != 0) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams(tail_res_row_num, __v * 2, 0, + __v * (q_heads - 1) * 2, 0)); + } +} + +// 平台函数：Ring LSE copyout（段落C 使用） +// 封装 ln_v + PIPE_BARRIER + add_v + PIPE_BARRIER + conv_v + SET/WAIT_FLAG + ub_to_gm_align + SET/WAIT_FLAG +// 用于 Stage2FinalizeAndOutput IS_RING 路径的 LSE 输出 +__aicore__ __attribute__((always_inline)) inline void PlatformRingLSECopyout( + AscendC::LocalTensor lse32_ubuf_tensor, + AscendC::LocalTensor gl32_ubuf_tensor, + AscendC::LocalTensor gm32_ubuf_tensor, + AscendC::LocalTensor lse_conv_ubuf_tensor, + AscendC::GlobalTensor lse_gm_tensor, + uint32_t sub_m, + uint32_t head_loop, + uint64_t o_offset) +{ + uint32_t sub_m_d64 = (sub_m + 63) / 64; + uint32_t lenBurst = sizeof(OUT_DTYPE); + ln_v(lse32_ubuf_tensor, gl32_ubuf_tensor, + sub_m_d64, 1, 1, 8, 8); + PIPE_BARRIER(V); + add_v(lse32_ubuf_tensor, lse32_ubuf_tensor, + gm32_ubuf_tensor, sub_m_d64, 1, 1, 1, 8, 8, 8); + PIPE_BARRIER(V); + conv_v(lse_conv_ubuf_tensor, + lse32_ubuf_tensor, sub_m_d64, 1, 1, 4, 8); + SET_FLAG(V, MTE3, EVENT_ID1); + WAIT_FLAG(V, MTE3, EVENT_ID1); + ub_to_gm_align( + lse_gm_tensor[(int64_t)(o_offset / __k)], + lse_conv_ubuf_tensor, + 0, 1, lenBurst * sub_m * head_loop, 0, 0, 0, 0); + SET_FLAG(MTE3, V, EVENT_ID1); + WAIT_FLAG(MTE3, V, EVENT_ID1); +} + +// 平台函数：Tail copyout gl/gm/go（段落C Tail 使用） +// 封装 ub_to_gm(gl) + ub_to_gm(gm) + SET_FLAG/WAIT_FLAG(V,MTE3,E5) + ub_to_gm(go) +// 用于 TailStage2FinalizeAndOutput 输出 gl/gm/go 三个 GM buffer +__aicore__ __attribute__((always_inline)) inline void PlatformTailCopyout( + AscendC::GlobalTensor gl_gm_tensor, + AscendC::GlobalTensor gm_gm_tensor, + AscendC::GlobalTensor go_gm_tensor, + AscendC::LocalTensor gl32_ubuf_tensor, + AscendC::LocalTensor gm32_ubuf_tensor, + AscendC::LocalTensor go32_ubuf_tensor, + uint32_t sub_m, + uint32_t round_v) +{ + ub_to_gm( + gl_gm_tensor, gl32_ubuf_tensor, + 0, 1, sub_m / FLOAT_BLOCK_SIZE, 0, 0); + + ub_to_gm( + gm_gm_tensor, gm32_ubuf_tensor, + 0, 1, sub_m / FLOAT_BLOCK_SIZE, 0, 0); + + SET_FLAG(V, MTE3, EVENT_ID5); + WAIT_FLAG(V, MTE3, EVENT_ID5); + ub_to_gm( + go_gm_tensor, go32_ubuf_tensor, + 0, 1, sub_m * round_v / FLOAT_BLOCK_SIZE, 0, 0); +} + +// ==================== Stage1 / 通用数据搬运平台函数 ==================== + +// 平台函数：MTE2→V 通用同步（SET_FLAG + WAIT_FLAG, EVENT_ID0） +// 用于 LoadQKDataFP16: gm_to_ub 后通知 V 管道可读 +__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSync() +{ + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); +} + +// 平台函数：V→MTE3 同步（SET_FLAG + WAIT_FLAG, EVENT_ID0） +// 用于 QuantizeAndOutput: conv_v 后通知 MTE3 可输出 +__aicore__ __attribute__((always_inline)) inline void PlatformVToMte3Sync() +{ + SET_FLAG(V, MTE3, EVENT_ID0); + WAIT_FLAG(V, MTE3, EVENT_ID0); +} + +// 平台函数：V→MTE2 通知（SET_FLAG + PIPE_BARRIER） +// 用于 QuantizeAndOutput: ReduceSum 后通知 MTE2 可搬入下一轮 +__aicore__ __attribute__((always_inline)) inline void PlatformVToMte2Notify() +{ + SET_FLAG(V, MTE2, EVENT_ID0); + PIPE_BARRIER(V); +} + +// 平台函数：V→MTE2 等待（WAIT_FLAG, EVENT_ID0） +// 用于 HeadLoop: 等待上一轮 V→MTE2 通知 +__aicore__ __attribute__((always_inline)) inline void PlatformVToMte2WaitEvent0() +{ + WAIT_FLAG(V, MTE2, EVENT_ID0); +} + +// 平台函数：MTE3→MTE2 等待（WAIT_FLAG, EVENT_ID4） +// 用于 HeadLoop: 等待 MTE3→MTE2 管道就绪 +__aicore__ __attribute__((always_inline)) inline void PlatformMte3ToMte2WaitEvent4() +{ + WAIT_FLAG(MTE3, MTE2, EVENT_ID4); +} + +// 平台函数：V→MTE2 通知（SET_FLAG, EVENT_ID0）— 仅 SET +// 用于 HeadLoop: 通知 MTE2 可读 go +__aicore__ __attribute__((always_inline)) inline void PlatformVToMte2SetEvent0() +{ + SET_FLAG(V, MTE2, EVENT_ID0); +} + +// 平台函数：MTE3→MTE2 通知（SET_FLAG, EVENT_ID4）— 仅 SET +// 用于 HeadLoop 尾部: 通知 MTE2 下一轮可搬入 +__aicore__ __attribute__((always_inline)) inline void PlatformMte3ToMte2SetEvent4() +{ + SET_FLAG(MTE3, MTE2, EVENT_ID4); +} + +// 平台函数：MTE3→MTE2 等待（WAIT_FLAG, EVENT_ID1）— TP1 专用 +// 用于 TP1 HeadLoop: 等待 MTE3→MTE2 EVENT_ID1 +__aicore__ __attribute__((always_inline)) inline void PlatformMte3ToMte2WaitEvent1() +{ + WAIT_FLAG(MTE3, MTE2, EVENT_ID1); +} + +// 平台函数：MTE3→MTE2 通知（SET_FLAG, EVENT_ID1）— TP1 专用 +// 用于 TP1 HeadLoop 尾部: 通知 MTE2 EVENT_ID1 +__aicore__ __attribute__((always_inline)) inline void PlatformMte3ToMte2SetEvent1() +{ + SET_FLAG(MTE3, MTE2, EVENT_ID1); +} + +// 平台函数：PIPE_BARRIER(ALL) +// 用于 Tail HeadLoop 尾部全管道屏障 +__aicore__ __attribute__((always_inline)) inline void PlatformPipeBarrierAll() +{ + PIPE_BARRIER(ALL); +} + +// ==================== 通用数据搬运平台函数 ==================== + +// 平台函数：gm_to_ub（通用 GM→UB 搬入，float 类型） +__aicore__ __attribute__((always_inline)) inline void PlatformGmToUbFloat( + AscendC::LocalTensor dst, + AscendC::GlobalTensor src, + uint32_t lenBurst) +{ + gm_to_ub(dst, src, 0, 1, lenBurst, 0, 0); +} + +// 平台函数：gm_to_ub（通用 GM→UB 搬入，模板类型） +template +__aicore__ __attribute__((always_inline)) inline void PlatformGmToUb( + AscendC::LocalTensor dst, + AscendC::GlobalTensor src, + uint32_t lenBurst) +{ + gm_to_ub(dst, src, 0, 1, lenBurst, 0, 0); +} + +// 平台函数：gm_to_ub（INT32 类型，CeilDiv 对齐） +__aicore__ __attribute__((always_inline)) inline void PlatformGmToUbInt32( + AscendC::LocalTensor dst, + AscendC::GlobalTensor src, + uint32_t sub_m, uint32_t qk_round_n) +{ + gm_to_ub( + dst, src, 0, 1, CeilDiv(sub_m * qk_round_n), 0, 0); +} + +// 平台函数：gm_to_ub_align（对齐 GM→UB，mmScaleType） +__aicore__ __attribute__((always_inline)) inline void PlatformGmToUbAlignMmScale( + AscendC::LocalTensor dst, + AscendC::GlobalTensor src, + uint32_t sub_m) +{ + gm_to_ub_align( + dst, src, 0, 1, sub_m * sizeof(mmScaleType), 0, 0, 0, 0); +} + +// 平台函数：ub_to_gm（通用 UB→GM 搬出，模板类型） +template +__aicore__ __attribute__((always_inline)) inline void PlatformUbToGm( + AscendC::GlobalTensor dst, + AscendC::LocalTensor src, + uint32_t lenBurst) +{ + ub_to_gm(dst, src, 0, 1, lenBurst, 0, 0); +} + +// 平台函数：ub_to_ub（通用 UB→UB 搬移，float 类型） +__aicore__ __attribute__((always_inline)) inline void PlatformUbToUb( + AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint32_t lenBurst) +{ + ub_to_ub(dst, src, 0, 1, lenBurst, 0, 0); +} + +// 平台函数：ub_to_ub（UB→UB，多 nBurst） +__aicore__ __attribute__((always_inline)) inline void PlatformUbToUbMulti( + AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint32_t nBurst, + uint32_t lenBurst, + uint32_t srcGap, + uint32_t dstGap) +{ + ub_to_ub(dst, src, 0, nBurst, lenBurst, srcGap, dstGap); +} + +// 平台函数：DataCopyPad（不规则 mask 加载） +__aicore__ __attribute__((always_inline)) inline void PlatformDataCopyPadMask( + AscendC::LocalTensor dst, + AscendC::GlobalTensor src, + uint32_t cur_q_seqlen, + uint32_t qk_n, + uint32_t maxKVSeqLen, + uint32_t qk_round_n) +{ + uint32_t aligned_mask_copy_len = RoundUp(qk_n); + uint32_t mask_dst_stride = (qk_round_n - aligned_mask_copy_len) / BLOCK_SIZE; + AscendC::DataCopyPad( + dst, src, + AscendC::DataCopyExtParams(cur_q_seqlen, qk_n * 2, maxKVSeqLen * 2 - qk_n * 2, mask_dst_stride, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0) + ); +} + +// 平台函数：DataCopy（规则 mask 加载） +__aicore__ __attribute__((always_inline)) inline void PlatformDataCopyMask( + AscendC::LocalTensor dst, + AscendC::GlobalTensor src, + uint32_t cur_q_seqlen, + uint32_t qk_round_n, + uint32_t maxKVSeqLen) +{ + AscendC::DataCopy( + dst, src, + AscendC::DataCopyParams( + cur_q_seqlen, + qk_round_n * 2 / 32, + MASK_COLUMNS * 2 / 32 - qk_round_n * 2 / 32, + 0) + ); +} + +// 平台函数：Cast mask→float +__aicore__ __attribute__((always_inline)) inline void PlatformCastMaskToFloat( + AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint32_t count) +{ + AscendC::Cast(dst, src, AscendC::RoundMode::CAST_NONE, count); +} + +// 平台函数：SetVectorMask 全置 -1 +__aicore__ __attribute__((always_inline)) inline void PlatformSetVectorMaskFull() +{ + SetVectorMask((uint64_t)-1, (uint64_t)-1); +} + +// 平台函数：__set_mask 设置部分 mask +__aicore__ __attribute__((always_inline)) inline void PlatformSetMask(uint32_t mask) +{ + __set_mask(mask); +} + +// ==================== 向量计算平台函数 ==================== + +// 平台函数：muls_v（标量乘法，循环 + tail mask） +__aicore__ __attribute__((always_inline)) inline void PlatformMulsVLoop( + AscendC::LocalTensor dst, + AscendC::LocalTensor src, + float scalar, + uint32_t sub_m, + uint32_t qk_n, + uint32_t qk_round_n) +{ + for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { + muls_v( + dst[vadd_idx * FLOAT_VECTOR_SIZE], + src[vadd_idx * FLOAT_VECTOR_SIZE], + scalar, sub_m, 1, 1, + qk_round_n / FLOAT_BLOCK_SIZE, + qk_round_n / FLOAT_BLOCK_SIZE); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + muls_v( + dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + scalar, sub_m, 1, 1, + qk_round_n / FLOAT_BLOCK_SIZE, + qk_round_n / FLOAT_BLOCK_SIZE); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); +} + +// 平台函数：Add（AscendC::Add，FP32 逐元素加） +__aicore__ __attribute__((always_inline)) inline void PlatformAddFloat( + AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint32_t count) +{ + AscendC::Add(dst, src0, src1, count); +} + +// 平台函数：max_v（逐元素取最大值） +__aicore__ __attribute__((always_inline)) inline void PlatformMaxV( + AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint32_t repeat) +{ + max_v( + dst, src0, src1, repeat, 1, 1, 1, 8, 8, 8); +} + +// 平台函数：sub_v（逐元素减法） +__aicore__ __attribute__((always_inline)) inline void PlatformSubV( + AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint32_t repeat) +{ + sub_v( + dst, src0, src1, repeat, 1, 1, 1, 8, 8, 8); +} + +// 平台函数：exp_v（指数运算） +__aicore__ __attribute__((always_inline)) inline void PlatformExpV( + AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint32_t repeat) +{ + exp_v( + dst, src, repeat, 1, 1, 8, 8); +} + +// 平台函数：mul_v（逐元素乘法） +__aicore__ __attribute__((always_inline)) inline void PlatformMulV( + AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint32_t repeat) +{ + mul_v( + dst, src0, src1, repeat, 1, 1, 1, 8, 8, 8); +} + +// 平台函数：add_v（逐元素加法） +__aicore__ __attribute__((always_inline)) inline void PlatformAddV( + AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint32_t repeat) +{ + add_v( + dst, src0, src1, repeat, 1, 1, 1, 8, 8, 8); +} + +// 平台函数：muls_v（标量乘法，单次调用） +__aicore__ __attribute__((always_inline)) inline void PlatformMulsV( + AscendC::LocalTensor dst, + AscendC::LocalTensor src, + float scalar, + uint32_t repeat) +{ + muls_v( + dst, src, scalar, repeat, 1, 1, 8, 8); +} + +// 平台函数：brcb_v（广播，uint32_t） +__aicore__ __attribute__((always_inline)) inline void PlatformBrcbV( + AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint32_t repeat) +{ + brcb_v(dst, src, 1, 8, repeat); +} + +// 平台函数：brcb_v（广播，float） +__aicore__ __attribute__((always_inline)) inline void PlatformBrcbVFloat( + AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint32_t round_sub_m) +{ + brcb_v( + dst.ReinterpretCast(), + src.ReinterpretCast(), + 1, 8, round_sub_m / FLOAT_BLOCK_SIZE); +} + +// 平台函数：conv_v（float→OUT_DTYPE 转换） +__aicore__ __attribute__((always_inline)) inline void PlatformConvVToOutput( + AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint32_t repeat) +{ + conv_v( + dst, src, repeat, 1, 1, 4, 8); +} + +// 平台函数：conv_v（int32_t→float 转换，含 repeat_times<255 分支） +__aicore__ __attribute__((always_inline)) inline void PlatformConvInt32ToFloat( + AscendC::LocalTensor dst, + AscendC::LocalTensor temp, + uint32_t count) +{ + uint32_t repeat_times = (count + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE; + if (repeat_times < 255) { + conv_v( + dst, temp, repeat_times, 1, 1, 8, 8); + } else { + for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { + conv_v( + dst[vconv_idx * count / 2], temp[vconv_idx * count / 2], + (count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + 1, 1, 8, 8); + } + } +} + +// 平台函数：conv_v（float→half 转换，含 repeat_times<255 分支） +__aicore__ __attribute__((always_inline)) inline void PlatformConvFloatToHalf( + AscendC::LocalTensor dst, + uint32_t count) +{ + uint32_t repeat_times = (count + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE; + if (repeat_times < 255) { + conv_v( + dst.template ReinterpretCast(), + dst.template ReinterpretCast(), + repeat_times, 1, 1, 4, 8); + } else { + for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { + conv_v( + dst.template ReinterpretCast()[vconv_idx * count / 2], + dst.template ReinterpretCast()[vconv_idx * count / 2], + (count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + 1, 1, 4, 8); + } + } + PIPE_BARRIER(V); +} + +// 平台函数：Cast 逐行转换（循环 + tail mask） +__aicore__ __attribute__((always_inline)) inline void PlatformCastHalfToInt8( + AscendC::LocalTensor dst, + uint32_t qk_n, + uint32_t qk_round_n, + uint32_t sub_m) +{ + for (uint32_t row_idx = 0; row_idx < qk_n / HALF_VECTOR_SIZE; ++row_idx) { + AscendC::Cast( + dst.template ReinterpretCast()[row_idx * HALF_VECTOR_SIZE], + dst.template ReinterpretCast()[row_idx * HALF_VECTOR_SIZE], + AscendC::RoundMode::CAST_RINT, + (uint64_t)0, sub_m, + {1, 1, (uint8_t)((qk_round_n) / BLOCK_SIZE), (uint8_t)(qk_round_n / BLOCK_SIZE)}); + } + if (qk_n % HALF_VECTOR_SIZE > 0) { + __set_mask(qk_n % HALF_VECTOR_SIZE); + AscendC::Cast( + dst.template ReinterpretCast()[qk_n / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], + dst.template ReinterpretCast()[qk_n / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], + AscendC::RoundMode::CAST_RINT, + (uint64_t)0, sub_m, + {1, 1, (uint8_t)((qk_round_n) / BLOCK_SIZE), (uint8_t)(qk_round_n / BLOCK_SIZE)}); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); +} + +// ==================== Tensor/Reduce 系列平台函数 ==================== +// 以下函数从 bs 层整体搬入，它们是平台原语+循环的封装，属于平台操作 + +// 平台函数：Tensor 逐行减法（brcb + sub_v 循环 + tail mask） +__aicore__ __attribute__((always_inline)) inline void PlatformTensorSubValueRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& MaxTensor, + const AscendC::LocalTensor& tempMaxTensor, + uint32_t sub_m, + uint32_t round_sub_m, + uint32_t qk_n, + uint32_t qk_round_n) +{ + brcb_v( + tempMaxTensor.ReinterpretCast(), + MaxTensor.ReinterpretCast(), + 1, 8, round_sub_m / FLOAT_BLOCK_SIZE); + PIPE_BARRIER(V); + for (uint32_t sub_v_idx = 0; sub_v_idx < qk_n / FLOAT_VECTOR_SIZE; ++sub_v_idx) { + sub_v( + dst[sub_v_idx * FLOAT_VECTOR_SIZE], + src[sub_v_idx * FLOAT_VECTOR_SIZE], + tempMaxTensor, sub_m, 1, 1, 0, + qk_round_n / FLOAT_BLOCK_SIZE, + qk_round_n / FLOAT_BLOCK_SIZE, 1); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + sub_v( + dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tempMaxTensor, sub_m, 1, 1, 0, + qk_round_n / FLOAT_BLOCK_SIZE, + qk_round_n / FLOAT_BLOCK_SIZE, 1); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); +} + +// 平台函数：Tensor 逐行除法（div_v 循环 + tail mask） +__aicore__ __attribute__((always_inline)) inline void PlatformTensorDivRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& src1, + uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n) +{ + PIPE_BARRIER(V); + for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { + div_v( + dst[vadd_idx * FLOAT_VECTOR_SIZE], + src[vadd_idx * FLOAT_VECTOR_SIZE], + src1, sub_m, 1, 1, 0, + qk_round_n / FLOAT_BLOCK_SIZE, + qk_round_n / FLOAT_BLOCK_SIZE, 1); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + div_v( + dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src1, sub_m, 1, 1, 0, + qk_round_n / FLOAT_BLOCK_SIZE, + qk_round_n / FLOAT_BLOCK_SIZE, 1); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); +} + +// 平台函数：Tensor 行归约求最大值（cmax_v + max_v 循环 + tail mask） +__aicore__ __attribute__((always_inline)) inline void PlatformReduceMaxRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& tempTensor, + uint32_t sub_m, + uint32_t qk_n, + uint32_t qk_round_n) +{ + if (qk_n <= FLOAT_VECTOR_SIZE) { + __set_mask(qk_n); + cmax_v( + dst, src, sub_m, 1, 1, qk_round_n / FLOAT_BLOCK_SIZE); + } else { + ub_to_ub( + tempTensor, src, 0, sub_m, + HALF_VECTOR_SIZE / BLOCK_SIZE, + (qk_round_n - FLOAT_VECTOR_SIZE) / FLOAT_BLOCK_SIZE, 0); + PIPE_BARRIER(V); + for (uint32_t rowmax_idx = 1; rowmax_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowmax_idx) { + max_v( + tempTensor, tempTensor, src[rowmax_idx * FLOAT_VECTOR_SIZE], + sub_m, 1, 1, 1, 8, 8, qk_round_n / FLOAT_BLOCK_SIZE); + PIPE_BARRIER(V); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + max_v( + tempTensor, tempTensor, + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + sub_m, 1, 1, 1, 8, 8, qk_round_n / FLOAT_BLOCK_SIZE); + } + PIPE_BARRIER(V); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + cmax_v( + dst, tempTensor, sub_m, 1, 1, 8); + } + SetVectorMask((uint64_t)-1, (uint64_t)-1); + PIPE_BARRIER(V); +} + +// 平台函数：Tensor 行归约求和（cadd_v + add_v 循环 + tail mask） +__aicore__ __attribute__((always_inline)) inline void PlatformReduceSumRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + uint32_t sub_m, + uint32_t qk_n, + uint32_t qk_round_n) +{ + if (qk_n <= FLOAT_VECTOR_SIZE) { + __set_mask(qk_n); + cadd_v( + dst, src, sub_m, 1, 1, qk_round_n / FLOAT_BLOCK_SIZE); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else { + for (uint32_t rowsum_idx = 1; rowsum_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowsum_idx) { + add_v( + src, src, src[rowsum_idx * FLOAT_VECTOR_SIZE], + sub_m, 1, 1, 1, + qk_round_n / FLOAT_BLOCK_SIZE, + qk_round_n / FLOAT_BLOCK_SIZE, + qk_round_n / FLOAT_BLOCK_SIZE); + PIPE_BARRIER(V); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + add_v( + src, src, + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + sub_m, 1, 1, 1, + qk_round_n / FLOAT_BLOCK_SIZE, + qk_round_n / FLOAT_BLOCK_SIZE, + qk_round_n / FLOAT_BLOCK_SIZE); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + cadd_v( + dst, src, sub_m, 1, 1, qk_round_n / FLOAT_BLOCK_SIZE); + } +} + +// 平台函数：Tensor 逐行乘法（mul_v 循环 + tail mask） +__aicore__ __attribute__((always_inline)) inline void PlatformTensorMulRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& src1, + uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t src1BlockStride) +{ + PIPE_BARRIER(V); + for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { + mul_v( + dst[vadd_idx * FLOAT_VECTOR_SIZE], + src[vadd_idx * FLOAT_VECTOR_SIZE], + src1, sub_m, 1, 1, src1BlockStride, + qk_round_n / FLOAT_BLOCK_SIZE, + qk_round_n / FLOAT_BLOCK_SIZE, 1); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + mul_v( + dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src1, sub_m, 1, 1, src1BlockStride, + qk_round_n / FLOAT_BLOCK_SIZE, + qk_round_n / FLOAT_BLOCK_SIZE, 1); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); +} + +// ==================== DeQuant/Quant 系列平台函数 ==================== + +// 平台函数：加载 deScale + online 乘 quantScale +__aicore__ __attribute__((always_inline)) inline void PlatformLoadDeScaleAndOnlineMul( + const AscendC::GlobalTensor& deScaleGm, + AscendC::LocalTensor deScaleUb, + AscendC::LocalTensor quantScale, + uint32_t sub_m, bool online) +{ + gm_to_ub_align( + deScaleUb, deScaleGm, 0, 1, sub_m * sizeof(mmScaleType), 0, 0, 0, 0); + if (online) { + PlatformMte2ToVSyncEvent2(); + PlatformTensorMulRepeatM(deScaleUb, deScaleUb, quantScale, 1, sub_m, RoundUp<16>(sub_m), 1); + } +} + +// 平台函数：加载 src(int32) + brcb 广播 deScale → tempScale +__aicore__ __attribute__((always_inline)) inline void PlatformLoadSrcAndBrcbScale( + const AscendC::GlobalTensor& src, + AscendC::LocalTensor temp, + AscendC::LocalTensor deScaleUb, + AscendC::LocalTensor tempScale, + uint32_t sub_m, uint32_t qk_round_n, bool move_tensor) +{ + if (move_tensor) { + gm_to_ub( + temp, src, 0, 1, CeilDiv(sub_m * qk_round_n), 0, 0); + } + PlatformMte2ToVSyncEvent0(); + brcb_v( + tempScale.template ReinterpretCast(), + deScaleUb.template ReinterpretCast(), + 1, 8, RoundUp<16>(sub_m) / FLOAT_BLOCK_SIZE); + PlatformVPipeBarrier(); +} + +// 平台函数：INT32→FP32 转换(conv_v) + 乘 tempScale +__aicore__ __attribute__((always_inline)) inline void PlatformConvInt32ToFP32AndMul( + AscendC::LocalTensor dst, + AscendC::LocalTensor temp, + AscendC::LocalTensor tempScale, + uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n) +{ + uint32_t count = sub_m * qk_round_n; + PlatformConvInt32ToFloat(dst, temp, count); + PlatformTensorMulRepeatM(dst, dst, tempScale, sub_m, qk_n, qk_round_n, 0); + PlatformVPipeBarrier(); +} + +// 平台函数：逐 Head 反量化（编排：加载+转换+乘法） +__aicore__ __attribute__((always_inline)) inline void PlatformDeQuantPerHead( + const AscendC::GlobalTensor& deScaleGm, + const AscendC::GlobalTensor& src, + AscendC::LocalTensor dst, + AscendC::LocalTensor temp, + AscendC::LocalTensor deScaleUb, + AscendC::LocalTensor tempScale, + AscendC::LocalTensor quantScale, + uint32_t sub_m, + uint32_t qk_n, + uint32_t qk_round_n, + bool online, + bool move_tensor) +{ + PlatformLoadDeScaleAndOnlineMul(deScaleGm, deScaleUb, quantScale, sub_m, online); + PlatformLoadSrcAndBrcbScale(src, temp, deScaleUb, tempScale, sub_m, qk_round_n, move_tensor); + PlatformConvInt32ToFP32AndMul(dst, temp, tempScale, sub_m, qk_n, qk_round_n); +} + +// 平台函数：逐 Token 量化（FP32 → FP16 → INT8） +__aicore__ __attribute__((always_inline)) inline void PlatformQuantPerToken( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& scale, + uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t pQuantOnline) +{ + if (pQuantOnline) { + PlatformTensorDivRepeatM(dst.template ReinterpretCast(), src, scale, sub_m, qk_n, qk_round_n); + } else { + PlatformTensorMulRepeatM(dst.template ReinterpretCast(), src, scale, sub_m, qk_n, qk_round_n, 0); + } + uint32_t count = sub_m * qk_round_n; + PlatformConvFloatToHalf(dst, count); + PlatformCastHalfToInt8(dst, qk_n, qk_round_n, sub_m); } \ No newline at end of file diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h index f307f14..2c77ad7 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h @@ -194,22 +194,16 @@ __aicore__ __attribute__((always_inline)) inline void LoadQKDataInt8( const uint32_t head_idx) { PlatformVToMte2Wait(); - DeQuantPerHeadImpl( - deq_scale_gm_tensor_q1[head_idx], - s_gm_tensor, + PlatformDeQuantPerHead( + deq_scale_gm_tensor_q1[head_idx], s_gm_tensor, ls32_quant_ubuf_tensor, ls32_quant_ubuf_tensor.template ReinterpretCast(), descale_q1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, qk_n, qk_round_n, 0, 1); - gm_to_ub( + PlatformGmToUbFloat( ls32_ubuf_tensor.template ReinterpretCast(), s_rope_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * qk_round_n / FLOAT_BLOCK_SIZE, - 0, // srcGap - 0 // dstGap - ); + sub_m * qk_round_n / FLOAT_BLOCK_SIZE); PlatformMte2ToVSync(); - AscendC::Add(ls32_ubuf_tensor, ls32_ubuf_tensor, ls32_quant_ubuf_tensor, sub_m * qk_round_n); // float + PlatformAddFloat(ls32_ubuf_tensor, ls32_ubuf_tensor, ls32_quant_ubuf_tensor, sub_m * qk_round_n); PlatformVPipeBarrier(); } @@ -226,52 +220,24 @@ __aicore__ __attribute__((always_inline)) inline void LoadQKDataFP16( bool need_mask) { PlatformVToMte2Wait(); - gm_to_ub( + PlatformGmToUb( ls32_ubuf_tensor.template ReinterpretCast(), s_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * qk_round_n / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); + sub_m * qk_round_n / FLOAT_BLOCK_SIZE); // mask 加载：mask_type==3 不规则 padding, mask_type==4 规则, 默认不加载 if (mask_type == 3) { - uint32_t aligned_mask_copy_len = RoundUp(qk_n); // 16 - uint32_t mask_dst_stride = (qk_round_n - aligned_mask_copy_len) / BLOCK_SIZE; // 0 - - AscendC::DataCopyPad( - mask_ubuf_tensor, - mask_gm_tensor, - AscendC::DataCopyExtParams( - cur_q_seqlen, - qk_n * 2, - maxKVSeqLen * 2 - qk_n * 2, - mask_dst_stride, - 0), - AscendC::DataCopyPadExtParams(false, 0, 0, 0) - ); + PlatformDataCopyPadMask(mask_ubuf_tensor, mask_gm_tensor, + cur_q_seqlen, qk_n, maxKVSeqLen, qk_round_n); } else if (need_mask && mask_type == 4) { - AscendC::DataCopy( - mask_ubuf_tensor, - mask_gm_tensor, - AscendC::DataCopyParams( - cur_q_seqlen, // blockCount - qk_round_n * 2 / 32, // blockLen, 2 is sizeof(half) - MASK_COLUMNS * 2 / 32 - qk_round_n * 2 / 32, // srcStride - 0 // dstStride - ) - ); + PlatformDataCopyMask(mask_ubuf_tensor, mask_gm_tensor, + cur_q_seqlen, qk_round_n, maxKVSeqLen); } PlatformMte2ToVSync(); if (mask_type == 3 || (need_mask && mask_type == 4)) { - AscendC::Cast( - mask32_ubuf_tensor, - mask_ubuf_tensor, - AscendC::RoundMode::CAST_NONE, + PlatformCastMaskToFloat(mask32_ubuf_tensor, mask_ubuf_tensor, cur_q_seqlen * qk_round_n); } } @@ -286,42 +252,17 @@ __aicore__ __attribute__((always_inline)) inline void ScaleAndMask( bool need_mask ) { - for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { - muls_v(ls32_ubuf_tensor[vadd_idx * FLOAT_VECTOR_SIZE], - ls32_ubuf_tensor[vadd_idx * FLOAT_VECTOR_SIZE], - tor, - sub_m, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride - ); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - muls_v(ls32_ubuf_tensor[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - ls32_ubuf_tensor[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tor, - sub_m, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); + PlatformMulsVLoop(ls32_ubuf_tensor, ls32_ubuf_tensor, tor, sub_m, qk_n, qk_round_n); if constexpr (tilingKeyType != TilingKeyType::TILING_INT8_DATA) { if (mask_type == 3 || (need_mask && mask_type == 4)) { uint32_t cur_compute_head_num = sub_m / cur_q_seqlen; for (uint32_t i = 0; i < cur_compute_head_num; i++) { - Add( + PlatformAddFloat( ls32_ubuf_tensor[cur_q_seqlen * qk_round_n * i], ls32_ubuf_tensor[cur_q_seqlen * qk_round_n * i], mask32_ubuf_tensor, - cur_q_seqlen * qk_round_n - ); + cur_q_seqlen * qk_round_n); } PlatformVPipeBarrier(); } @@ -342,58 +283,22 @@ __aicore__ __attribute__((always_inline)) inline void UpdateSoftmaxState( uint32_t round_sub_m = (sub_m + 15) / 16 * 16; // *** lm = rowmax(ls) - ReduceMaxRepeatM(lm32_ubuf_tensor, ls32_ubuf_tensor, lp32_ubuf_tensor, sub_m, qk_n, qk_round_n); + PlatformReduceMaxRepeatM(lm32_ubuf_tensor, ls32_ubuf_tensor, lp32_ubuf_tensor, sub_m, qk_n, qk_round_n); if (n_idx != 0) { // *** hm = vmax(lm, gm) - max_v(hm32_ubuf_tensor, - lm32_ubuf_tensor, - gm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); + PlatformMaxV(hm32_ubuf_tensor, lm32_ubuf_tensor, gm32_ubuf_tensor, sub_m_d64); PlatformVPipeBarrier(); // *** dm = gm - hm - sub_v(dm32_ubuf_tensor, - gm32_ubuf_tensor, - hm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); + PlatformSubV(dm32_ubuf_tensor, gm32_ubuf_tensor, hm32_ubuf_tensor, sub_m_d64); PlatformVPipeBarrier(); } else { // *** hm = lm - ub_to_ub( - hm32_ubuf_tensor, - lm32_ubuf_tensor, - 0, // sid - 1, // nBurst - round_sub_m / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); + PlatformUbToUb(hm32_ubuf_tensor, lm32_ubuf_tensor, round_sub_m / FLOAT_BLOCK_SIZE); PlatformVPipeBarrier(); } // *** gm = hm - ub_to_ub( - gm32_ubuf_tensor, - hm32_ubuf_tensor, - 0, // sid - 1, // nBurst - round_sub_m / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - PIPE_BARRIER(V); + PlatformUbToUb(gm32_ubuf_tensor, hm32_ubuf_tensor, round_sub_m / FLOAT_BLOCK_SIZE); + PlatformVPipeBarrier(); } // 业务子函数4：减法 + 指数 @@ -407,19 +312,13 @@ __aicore__ __attribute__((always_inline)) inline void SubAndExp( uint32_t round_sub_m = (sub_m + 15) / 16 * 16; // *** ls = ls - hm_block - TensorSubValueRepeatM(ls32_ubuf_tensor, ls32_ubuf_tensor, + PlatformTensorSubValueRepeatM(ls32_ubuf_tensor, ls32_ubuf_tensor, hm32_ubuf_tensor, tv32_ubuf_tensor, sub_m, round_sub_m, qk_n, qk_round_n); // *** ls = exp(ls) - exp_v(ls32_ubuf_tensor, - ls32_ubuf_tensor, - (sub_m * qk_round_n + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - PIPE_BARRIER(V); + PlatformExpV(ls32_ubuf_tensor, ls32_ubuf_tensor, + (sub_m * qk_round_n + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE); + PlatformVPipeBarrier(); } // 业务子函数5：量化/转换输出 @@ -440,71 +339,28 @@ __aicore__ __attribute__((always_inline)) inline void QuantizeAndOutput( float quantMax = (float)1 / (float)127; if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - sub_v(pm32_ubuf_tensor, - lm32_ubuf_tensor, - hm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); + PlatformSubV(pm32_ubuf_tensor, lm32_ubuf_tensor, hm32_ubuf_tensor, sub_m_d64); PlatformVPipeBarrier(); - exp_v(pm32_ubuf_tensor, - pm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); + PlatformExpV(pm32_ubuf_tensor, pm32_ubuf_tensor, sub_m_d64); PlatformVPipeBarrier(); - muls_v(pm32_ubuf_tensor, - pm32_ubuf_tensor, - quantMax, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); + PlatformMulsV(pm32_ubuf_tensor, pm32_ubuf_tensor, quantMax, sub_m_d64); PlatformVPipeBarrier(); - brcb_v( - tv32_ubuf_tensor.ReinterpretCast(), - pm32_ubuf_tensor.ReinterpretCast(), - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - QuantPerTokenImpl(lp_ubuf_tensor, ls32_ubuf_tensor, tv32_ubuf_tensor, sub_m, qk_n, qk_round_n, 1); + PlatformBrcbV(tv32_ubuf_tensor.template ReinterpretCast(), + pm32_ubuf_tensor.template ReinterpretCast(), round_sub_m / FLOAT_BLOCK_SIZE); + PlatformQuantPerToken(lp_ubuf_tensor, ls32_ubuf_tensor, tv32_ubuf_tensor, sub_m, qk_n, qk_round_n, 1); } else { - conv_v(lp_ubuf_tensor, - ls32_ubuf_tensor, - (sub_m * qk_round_n + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); + PlatformConvVToOutput(lp_ubuf_tensor, ls32_ubuf_tensor, + (sub_m * qk_round_n + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE); PlatformVPipeBarrier(); } PlatformVToMte3Sync(); - ub_to_gm( - p_gm_tensor, - lp_ubuf_tensor, - 0, // sid - 1, // nBurst - sub_m * qk_round_n * T_BLOCK_OFFSET / T_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); + PlatformUbToGm(p_gm_tensor, lp_ubuf_tensor, + sub_m * qk_round_n * T_BLOCK_OFFSET / T_BLOCK_SIZE); // *** ll = rowsum(ls32) - ReduceSumRepeatM(ll_ubuf_tensor, ls32_ubuf_tensor, sub_m, qk_n, qk_round_n); + PlatformReduceSumRepeatM(ll_ubuf_tensor, ls32_ubuf_tensor, sub_m, qk_n, qk_round_n); PlatformVToMte2Notify(); - PIPE_BARRIER(V); + PlatformVPipeBarrier(); } // 业务函数：非TP1 Phase 1 — SoftmaxStage1 编排（调用5个业务子函数） @@ -631,6 +487,110 @@ __aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage1( PlatformSoftmaxStage1PostSync(); } +// ==================== Stage2 业务子函数（三层拆分）==================== +// 以下子函数提取自 SoftmaxStage2MLAHeadLoop / TP1 / Tail 三个函数的共性段落， +// 通过平台原语下沉消除重复代码。 + +// 业务子函数：段落B — Merge Accumulate（brcb dm + load go + go=go*dm + go=lo+go） +// 非TP1/TP1/Tail 三函数完全一致的段落B逻辑 +__aicore__ __attribute__((always_inline)) inline void Stage2MergeAccumulate( + AscendC::GlobalTensor go_gm_tensor, + AscendC::LocalTensor go32_ubuf_tensor, + AscendC::LocalTensor lo_ubuf_tensor, + AscendC::LocalTensor tv32_ubuf_tensor, + AscendC::LocalTensor dm32_ubuf_tensor, + uint32_t sub_m, + uint32_t round_sub_m, + uint32_t round_v, + uint32_t head_loop) +{ + // brcb dm → tv, go = go * dm_block + PlatformMulVectorByBroadcast(go32_ubuf_tensor, tv32_ubuf_tensor, + dm32_ubuf_tensor, sub_m, round_sub_m, round_v, __v); + + if (head_loop > 1) { + PlatformLoadGoFromGM(go_gm_tensor, go32_ubuf_tensor, sub_m, round_v); + } + + // go = lo + go + PlatformAddLoToGo(go32_ubuf_tensor, lo_ubuf_tensor, sub_m, round_v); +} + +// 业务子函数：段落D — Intermediate Save（保存 go32 到 GM，非最后一轮） +// 非TP1/TP1/Tail 三函数完全一致的段落D逻辑 +__aicore__ __attribute__((always_inline)) inline void Stage2IntermediateSave( + AscendC::GlobalTensor go_gm_tensor, + AscendC::LocalTensor go32_ubuf_tensor, + uint32_t sub_m, + uint32_t round_v) +{ + PlatformIntermediateSave(go_gm_tensor, go32_ubuf_tensor, sub_m, round_v); +} + +// 业务子函数：段落C — Finalize & Output（div + conv + DataCopyPad + Ring LSE） +// 非TP1/TP1 共用的最终输出逻辑（Tail 不使用此函数） +__aicore__ __attribute__((always_inline)) inline void Stage2FinalizeAndOutput( + AscendC::GlobalTensor o_gm_tensor, + AscendC::GlobalTensor go_gm_tensor, + AscendC::LocalTensor go32_ubuf_tensor, + AscendC::LocalTensor go_ubuf_tensor, + AscendC::LocalTensor tv32_ubuf_tensor, + AscendC::LocalTensor gl32_ubuf_tensor, + AscendC::LocalTensor lse32_ubuf_tensor, + AscendC::LocalTensor lse_conv_ubuf_tensor, + AscendC::LocalTensor gm32_ubuf_tensor, + AscendC::GlobalTensor lse_gm_tensor, + uint32_t sub_m, + uint32_t round_sub_m, + uint32_t round_v, + uint32_t head_loop, + uint32_t head_loop_idx, + uint32_t q_seq_len, + uint32_t numhead_per_process, + uint32_t head_res_row_num, + uint32_t head_start_sblock_idx, + uint32_t tail_res_row_num, + uint64_t o_offset) +{ + uint32_t sub_m_d64 = (sub_m + 63) / 64; + + // go = go / gl_block + PlatformDivVectorByBroadcast(go32_ubuf_tensor, tv32_ubuf_tensor, + gl32_ubuf_tensor, sub_m, round_sub_m, round_v, __v, head_loop_idx); + + // go = castfp32to16(go) + V→MTE3 同步 + PlatformConvGoToOutput(go_ubuf_tensor, go32_ubuf_tensor, sub_m, round_v); + + // DataCopyPad output + PlatformDataCopyPadOutput(o_gm_tensor, go_ubuf_tensor, + sub_m, round_v, q_seq_len, numhead_per_process, + head_res_row_num, head_start_sblock_idx, tail_res_row_num); + + // Ring LSE copyout + if constexpr (IS_RING) { + PlatformRingLSECopyout(lse32_ubuf_tensor, gl32_ubuf_tensor, gm32_ubuf_tensor, + lse_conv_ubuf_tensor, lse_gm_tensor, sub_m, head_loop, o_offset); + } +} + +// 业务子函数：段落C — Tail Finalize & Output（仅 copyout gl/gm/go，无 div/conv） +// Tail 专用最终输出逻辑 +__aicore__ __attribute__((always_inline)) inline void TailStage2FinalizeAndOutput( + AscendC::GlobalTensor go_gm_tensor, + AscendC::GlobalTensor gl_gm_tensor, + AscendC::GlobalTensor gm_gm_tensor, + AscendC::LocalTensor go32_ubuf_tensor, + AscendC::LocalTensor gl32_ubuf_tensor, + AscendC::LocalTensor gm32_ubuf_tensor, + uint32_t sub_m, + uint32_t round_v) +{ + // Copyout gl/gm/go to GM + PlatformTailCopyout(gl_gm_tensor, gm_gm_tensor, go_gm_tensor, + gl32_ubuf_tensor, gm32_ubuf_tensor, go32_ubuf_tensor, + sub_m, round_v); +} + // 业务函数：非TP1 Phase 2 — SoftmaxStage2MLAHeadLoop 核心 Head Loop 计算 // 包含 load lo → DeQuant → exp(dm) → gl=dm*gl → gl=ll+gl → brcb dm → go=go*dm → go=lo+go // → brcb gl → go=go/gl → conv_v → DataCopyPad输出 → Ring LSE copyout @@ -662,337 +622,75 @@ __aicore__ __attribute__((always_inline)) inline void SoftmaxStage2MLAHeadLoop( { uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - WAIT_FLAG(V, MTE2, EVENT_ID0); + PlatformVToMte2WaitEvent0(); if (n_idx != 0) { - gm_to_ub( + PlatformGmToUb( lo_ubuf_tensor.template ReinterpretCast(), o_tmp_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); + sub_m * round_v / FLOAT_BLOCK_SIZE); + PlatformMte2ToVSync(); if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - DeQuantPerHeadImpl( + PlatformDeQuantPerHead( deq_scale_gm_tensor_k1[head_idx], o_tmp_gm_tensor, - lo_ubuf_tensor, lo_ubuf_tensor.template ReinterpretCast(),// lo_ubuf_tensor use the same ptr + lo_ubuf_tensor, lo_ubuf_tensor.template ReinterpretCast(), descale_k1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, round_v, round_v, 1, 0); } } - SetVectorMask((uint64_t)-1, (uint64_t)-1); - WAIT_FLAG(MTE3, MTE2, EVENT_ID4); + PlatformSetVectorMaskFull(); + PlatformMte3ToMte2WaitEvent4(); if (n_idx != 0) { // *** dm = exp(dm) if (head_loop_idx == 0) { - exp_v(dm32_ubuf_tensor, - dm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - PIPE_BARRIER(V); + PlatformExpV(dm32_ubuf_tensor, dm32_ubuf_tensor, sub_m_d64); + PlatformVPipeBarrier(); // *** gl = dm * gl - mul_v(gl32_ubuf_tensor, - dm32_ubuf_tensor, - gl32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); + PlatformMulV(gl32_ubuf_tensor, dm32_ubuf_tensor, gl32_ubuf_tensor, sub_m_d64); + PlatformVPipeBarrier(); // *** gl = ll + gl - add_v(gl32_ubuf_tensor, - gl32_ubuf_tensor, - ll_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); - } - SetVectorMask((uint64_t)-1, (uint64_t)-1); - brcb_v(tv32_ubuf_tensor.ReinterpretCast(), - dm32_ubuf_tensor.ReinterpretCast(), - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - PIPE_BARRIER(V); - if (head_loop > 1) { - gm_to_ub( - go32_ubuf_tensor, - go_gm_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - } - - // *** go = go * dm_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { - mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (__v % FLOAT_VECTOR_SIZE > 0) { - __set_mask(__v % FLOAT_VECTOR_SIZE); - mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); + PlatformAddV(gl32_ubuf_tensor, gl32_ubuf_tensor, ll_ubuf_tensor, sub_m_d64); + PlatformVPipeBarrier(); } - PIPE_BARRIER(V); - // *** go = lo + go - add_v(go32_ubuf_tensor, - go32_ubuf_tensor, - lo_ubuf_tensor, - (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); + // *** 段落B：brcb dm + load go + go=go*dm + go=lo+go + Stage2MergeAccumulate(go_gm_tensor, go32_ubuf_tensor, lo_ubuf_tensor, + tv32_ubuf_tensor, dm32_ubuf_tensor, + sub_m, round_sub_m, round_v, head_loop); } else { // *** gl = ll if (head_loop_idx == 0) { - ub_to_ub( - gl32_ubuf_tensor, - ll_ubuf_tensor, - 0, // sid - 1, // nBurst - 64 / FLOAT_BLOCK_SIZE, // lenBurst - // round_sub_m / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - PIPE_BARRIER(V); + PlatformUbToUb(gl32_ubuf_tensor, ll_ubuf_tensor, 64 / FLOAT_BLOCK_SIZE); + PlatformVPipeBarrier(); } - gm_to_ub( + PlatformGmToUb( go32_ubuf_tensor.template ReinterpretCast(), o_tmp_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); + sub_m * round_v / FLOAT_BLOCK_SIZE); if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - DeQuantPerHeadImpl( + PlatformDeQuantPerHead( deq_scale_gm_tensor_k1[head_idx], o_tmp_gm_tensor, go32_ubuf_tensor, go32_ubuf_tensor.template ReinterpretCast(), descale_k1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, round_v, round_v, 1, 0); } else { - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); + PlatformMte2ToVSync(); } } - SET_FLAG(V, MTE2, EVENT_ID0); + PlatformVToMte2SetEvent0(); if (n_idx == n_loop - 1) { - // *** gl_block = expand_to_block(gl) - brcb_v(tv32_ubuf_tensor.ReinterpretCast(), - gl32_ubuf_tensor.ReinterpretCast()[head_loop_idx * 16], - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - PIPE_BARRIER(V); - // *** go = go / gl_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vdiv_idx = 0; vdiv_idx < __v / FLOAT_VECTOR_SIZE; ++vdiv_idx) { - div_v(go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (__v % FLOAT_VECTOR_SIZE > 0) { - __set_mask(__v % FLOAT_VECTOR_SIZE); - div_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); // fix hidden_size=96 - } - PIPE_BARRIER(V); - - // *** go = castfp32to16(go) - conv_v(go_ubuf_tensor, - go32_ubuf_tensor, - (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); - SET_FLAG(V, MTE3, EVENT_ID0); - WAIT_FLAG(V, MTE3, EVENT_ID0); - - uint32_t inner_o_gm_offset = 0; - uint32_t inner_go_ubuf_offset = 0; - - if (head_res_row_num != 0) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset + q_heads * __v * head_start_sblock_idx], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams( - head_res_row_num, // blockCount - __v * 2, // blockLen - 0, // srcStride - __v * (q_heads - 1) * 2, // dstStride - 0 // rsv - ) - ); - inner_o_gm_offset += __v; - inner_go_ubuf_offset += head_res_row_num * __v; - } - - for (uint32_t i = 0; i < numhead_per_process; i++) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams( - q_seq_len, // blockCount - __v * 2, // blockLen - 0, // srcStride - __v * (q_heads - 1) * 2, // dstStride - 0 // rsv - ) - ); - inner_o_gm_offset += __v; - inner_go_ubuf_offset += q_seq_len * __v; - } - - if (tail_res_row_num != 0) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams( - tail_res_row_num, // blockCount - __v * 2, // blockLen - 0, // srcStride - __v * (q_heads - 1) * 2, // dstStride - 0 // rsv - ) - ); - } - // ********************* move O to GM ************************ - if constexpr (IS_RING) { - uint32_t lenBurst = sizeof(OUT_DTYPE); - ln_v(lse32_ubuf_tensor, - gl32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - PIPE_BARRIER(V); - add_v(lse32_ubuf_tensor, - lse32_ubuf_tensor, - gm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); - conv_v(lse_conv_ubuf_tensor, - lse32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); - SET_FLAG(V, MTE3, EVENT_ID1); - WAIT_FLAG(V, MTE3, EVENT_ID1); - // copyout lse - ub_to_gm_align( - lse_gm_tensor[(int64_t)(o_offset / __k)], - lse_conv_ubuf_tensor, - 0, // sid - 1, // nBurst - lenBurst * sub_m * head_loop, // lenBurst - 0, // leftPaddingNum - 0, // rightPaddingNum - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE3, V, EVENT_ID1); - WAIT_FLAG(MTE3, V, EVENT_ID1); - } - + // *** 段落C：brcb gl + go=go/gl + conv + DataCopyPad输出 + Ring LSE + Stage2FinalizeAndOutput(o_gm_tensor, go_gm_tensor, go32_ubuf_tensor, go_ubuf_tensor, + tv32_ubuf_tensor, gl32_ubuf_tensor, lse32_ubuf_tensor, lse_conv_ubuf_tensor, + gm32_ubuf_tensor, lse_gm_tensor, + sub_m, round_sub_m, round_v, head_loop, head_loop_idx, + q_seq_len, numhead_per_process, head_res_row_num, + head_start_sblock_idx, tail_res_row_num, o_offset); } else if (head_loop > 1) { - SET_FLAG(V, MTE3, EVENT_ID5); - WAIT_FLAG(V, MTE3, EVENT_ID5); - ub_to_gm( - go_gm_tensor, - go32_ubuf_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); + // *** 段落D：中间结果保存 SET_FLAG + WAIT + ub_to_gm go32 + Stage2IntermediateSave(go_gm_tensor, go32_ubuf_tensor, sub_m, round_v); } - SET_FLAG(MTE3, MTE2, EVENT_ID4); + PlatformMte3ToMte2SetEvent4(); } // 业务函数：非TP1 Phase 2 — Softmax Stage2 调度 @@ -1202,272 +900,46 @@ __aicore__ __attribute__((always_inline)) inline void SoftmaxStage2MLAHeadLoopTP { uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - WAIT_FLAG(V, MTE2, EVENT_ID0); + PlatformVToMte2WaitEvent0(); if (n_idx != 4) { - gm_to_ub( + PlatformGmToUb( lo_ubuf_tensor.template ReinterpretCast(), o_tmp_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); + sub_m * round_v / FLOAT_BLOCK_SIZE); + PlatformMte2ToVSync(); } - SetVectorMask((uint64_t)-1, (uint64_t)-1); - WAIT_FLAG(MTE3, MTE2, EVENT_ID4); + PlatformSetVectorMaskFull(); + PlatformMte3ToMte2WaitEvent4(); if (n_idx != 4) { // expand_to_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - brcb_v(tv32_ubuf_tensor.ReinterpretCast(), - dm32_ubuf_tensor.ReinterpretCast(), - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - PIPE_BARRIER(V); - if (head_loop > 1) { - gm_to_ub( - go32_ubuf_tensor, - go_gm_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - } - - // *** go = go * dm_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { - mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (__v % FLOAT_VECTOR_SIZE > 0) { - __set_mask(__v % FLOAT_VECTOR_SIZE); - mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); - // *** go = lo + go - add_v(go32_ubuf_tensor, - go32_ubuf_tensor, - lo_ubuf_tensor, - (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); + // *** 段落B：brcb dm + load go + go=go*dm + go=lo+go + Stage2MergeAccumulate(go_gm_tensor, go32_ubuf_tensor, lo_ubuf_tensor, + tv32_ubuf_tensor, dm32_ubuf_tensor, + sub_m, round_sub_m, round_v, head_loop); } else { // *** go = lo - gm_to_ub( + PlatformGmToUb( go32_ubuf_tensor.template ReinterpretCast(), o_tmp_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); + sub_m * round_v / FLOAT_BLOCK_SIZE); + PlatformMte2ToVSync(); } - SET_FLAG(V, MTE2, EVENT_ID0); + PlatformVToMte2SetEvent0(); if (n_idx + 4 > n_loop + 4 - 1) { - // *** gl_block = expand_to_block(gl), 存放于 tv - brcb_v(tv32_ubuf_tensor.ReinterpretCast(), - gl32_ubuf_tensor.ReinterpretCast()[head_loop_idx * 16], - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - PIPE_BARRIER(V); - // *** go = go / gl_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vdiv_idx = 0; vdiv_idx < __v / FLOAT_VECTOR_SIZE; ++vdiv_idx) { - div_v(go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (__v % FLOAT_VECTOR_SIZE > 0) { - __set_mask(__v % FLOAT_VECTOR_SIZE); - div_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); // fix hidden_size=96 - } - PIPE_BARRIER(V); - - // *** go = castfp32to16(go) - conv_v(go_ubuf_tensor, - go32_ubuf_tensor, - (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); - SET_FLAG(V, MTE3, EVENT_ID0); - WAIT_FLAG(V, MTE3, EVENT_ID0); - - uint32_t inner_o_gm_offset = 0; - uint32_t inner_go_ubuf_offset = 0; - - if (head_res_row_num != 0) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset + q_heads * __v * head_start_sblock_idx], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams( - head_res_row_num, // blockCount - __v * 2, // blockLen - 0, // srcStride - __v * (q_heads - 1) * 2, // dstStride - 0 // rsv - ) - ); - inner_o_gm_offset += __v; - inner_go_ubuf_offset += head_res_row_num * __v; - } - - for (uint32_t i = 0; i < numhead_per_process; i++) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams( - q_seq_len, // blockCount - __v * 2, // blockLen - 0, // srcStride - __v * (q_heads - 1) * 2, // dstStride - 0 // rsv - ) - ); - inner_o_gm_offset += __v; - inner_go_ubuf_offset += q_seq_len * __v; - } - - if (tail_res_row_num != 0) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams( - tail_res_row_num, // blockCount - __v * 2, // blockLen - 0, // srcStride - __v * (q_heads - 1) * 2, // dstStride - 0 // rsv - ) - ); - } - // ********************* move O to GM ************************ - if constexpr (IS_RING) { - uint32_t lenBurst = sizeof(OUT_DTYPE); - ln_v(lse32_ubuf_tensor, - gl32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - PIPE_BARRIER(V); - add_v(lse32_ubuf_tensor, - lse32_ubuf_tensor, - gm32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); - conv_v(lse_conv_ubuf_tensor, - lse32_ubuf_tensor, - sub_m_d64, // repeat - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); - SET_FLAG(V, MTE3, EVENT_ID1); - WAIT_FLAG(V, MTE3, EVENT_ID1); - // copyout lse - ub_to_gm_align( - lse_gm_tensor[(int64_t)(o_offset / __k)], - lse_conv_ubuf_tensor, - 0, // sid - 1, // nBurst - lenBurst * sub_m * head_loop, // lenBurst - 0, // leftPaddingNum - 0, // rightPaddingNum - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE3, V, EVENT_ID1); - WAIT_FLAG(MTE3, V, EVENT_ID1); - } - } - else if (head_loop > 1) { - SET_FLAG(V, MTE3, EVENT_ID5); - WAIT_FLAG(V, MTE3, EVENT_ID5); - ub_to_gm( - go_gm_tensor, - go32_ubuf_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); + // *** 段落C：brcb gl + go=go/gl + conv + DataCopyPad输出 + Ring LSE + Stage2FinalizeAndOutput(o_gm_tensor, go_gm_tensor, go32_ubuf_tensor, go_ubuf_tensor, + tv32_ubuf_tensor, gl32_ubuf_tensor, lse32_ubuf_tensor, lse_conv_ubuf_tensor, + gm32_ubuf_tensor, lse_gm_tensor, + sub_m, round_sub_m, round_v, head_loop, head_loop_idx, + q_seq_len, numhead_per_process, head_res_row_num, + head_start_sblock_idx, tail_res_row_num, o_offset); + } else if (head_loop > 1) { + // *** 段落D：中间结果保存 SET_FLAG + WAIT + ub_to_gm go32 + Stage2IntermediateSave(go_gm_tensor, go32_ubuf_tensor, sub_m, round_v); } - SET_FLAG(MTE3, MTE2, EVENT_ID4); + PlatformMte3ToMte2SetEvent4(); } // 业务函数：TP1 Phase 2 — SoftmaxStage2MLAHeadLoopTP1 调度（非 Tail 路径） @@ -1557,161 +1029,44 @@ __aicore__ __attribute__((always_inline)) inline void TailSoftmaxStage2MLAHeadLo { uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - WAIT_FLAG(V, MTE2, EVENT_ID0); + PlatformVToMte2WaitEvent0(); if (n_idx != 4) { - gm_to_ub( + PlatformGmToUb( lo_ubuf_tensor.template ReinterpretCast(), o_tmp_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); + sub_m * round_v / FLOAT_BLOCK_SIZE); + PlatformMte2ToVSync(); } - SetVectorMask((uint64_t)-1, (uint64_t)-1); - WAIT_FLAG(MTE3, MTE2, EVENT_ID4); + PlatformSetVectorMaskFull(); + PlatformMte3ToMte2WaitEvent4(); if (n_idx != 4) { // expand_to_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - brcb_v(tv32_ubuf_tensor.ReinterpretCast(), - dm32_ubuf_tensor.ReinterpretCast(), - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - PIPE_BARRIER(V); - if (head_loop > 1) { - gm_to_ub( - go32_ubuf_tensor, - go_gm_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); - } - - // *** go = go * dm_block - SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { - mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (__v % FLOAT_VECTOR_SIZE > 0) { - __set_mask(__v % FLOAT_VECTOR_SIZE); - mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); - // *** go = lo + go - add_v(go32_ubuf_tensor, - go32_ubuf_tensor, - lo_ubuf_tensor, - (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - 8 // src1RepeatStride - ); - PIPE_BARRIER(V); + // *** 段落B：brcb dm + load go + go=go*dm + go=lo+go + Stage2MergeAccumulate(go_gm_tensor, go32_ubuf_tensor, lo_ubuf_tensor, + tv32_ubuf_tensor, dm32_ubuf_tensor, + sub_m, round_sub_m, round_v, head_loop); } else { // *** go = lo - gm_to_ub( + PlatformGmToUb( go32_ubuf_tensor.template ReinterpretCast(), o_tmp_gm_tensor, - 0, // sid - 1, // nBurst - sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst - 0, // srcGap - 0 // dstGap - ); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); + sub_m * round_v / FLOAT_BLOCK_SIZE); + PlatformMte2ToVSync(); } - SET_FLAG(V, MTE2, EVENT_ID0); + PlatformVToMte2SetEvent0(); if (n_idx + 4 > n_loop + 4 - 1) { - // The last step to process the o with dividing and copyout - // TODO: Maybe the following two don't need waiting - // Copyout gl32_ubuf_tensor to gl_gm_tensor - ub_to_gm( - gl_gm_tensor, - gl32_ubuf_tensor, - 0, - 1, - sub_m / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - - // Copyout rowmax to global gm32_ubuf_tensor - ub_to_gm( - gm_gm_tensor, - gm32_ubuf_tensor, - 0, - 1, - sub_m / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - // Copyout go32_ubuf_tensor to go_gm_tensor - // This is needed to wait for former calculation - SET_FLAG(V, MTE3, EVENT_ID5); - WAIT_FLAG(V, MTE3, EVENT_ID5); - ub_to_gm( - go_gm_tensor, - go32_ubuf_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); - - } - else if (head_loop > 1) { - SET_FLAG(V, MTE3, EVENT_ID5); - WAIT_FLAG(V, MTE3, EVENT_ID5); - ub_to_gm( - go_gm_tensor, - go32_ubuf_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); + // *** 段落C（Tail）：copyout gl/gm/go + TailStage2FinalizeAndOutput(go_gm_tensor, gl_gm_tensor, gm_gm_tensor, + go32_ubuf_tensor, gl32_ubuf_tensor, gm32_ubuf_tensor, + sub_m, round_v); + } else if (head_loop > 1) { + // *** 段落D：中间结果保存 SET_FLAG + WAIT + ub_to_gm go32 + Stage2IntermediateSave(go_gm_tensor, go32_ubuf_tensor, sub_m, round_v); } - SET_FLAG(MTE3, MTE2, EVENT_ID4); - PIPE_BARRIER(ALL); + PlatformMte3ToMte2SetEvent4(); + PlatformPipeBarrierAll(); } // 业务函数：TP1 Phase 2 (Tail) — TailSoftmaxStage2MLAHeadLoopTP1 调度（Tail 路径） @@ -1824,7 +1179,7 @@ __aicore__ __attribute__((always_inline)) inline void TailInnerRunVectorChangeTP // ====== Tensor 逐行重复计算（Vector 业务基础函数）====== -// Tensor 逐行减法（brcb 广播 Max + sub_v 循环 + tail 处理） +// Tensor 逐行减法（委托 arch 层平台函数） __aicore__ __attribute__((always_inline)) inline void TensorSubValueRepeatM( const AscendC::LocalTensor& dst, const AscendC::LocalTensor& src, @@ -1835,85 +1190,20 @@ __aicore__ __attribute__((always_inline)) inline void TensorSubValueRepeatM( uint32_t qk_n, uint32_t qk_round_n) { - brcb_v( - tempMaxTensor.ReinterpretCast(), - MaxTensor.ReinterpretCast(), - 1, // dstBlockStride - 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat - ); - PIPE_BARRIER(V); - for (uint32_t sub_v_idx = 0; sub_v_idx < qk_n / FLOAT_VECTOR_SIZE; ++sub_v_idx) { - sub_v(dst[sub_v_idx * FLOAT_VECTOR_SIZE], - src[sub_v_idx * FLOAT_VECTOR_SIZE], - tempMaxTensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - sub_v(dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tempMaxTensor, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); + PlatformTensorSubValueRepeatM(dst, src, MaxTensor, tempMaxTensor, sub_m, round_sub_m, qk_n, qk_round_n); } -// Tensor 逐行除法（div_v 循环 + tail 处理） +// Tensor 逐行除法（委托 arch 层平台函数） __aicore__ __attribute__((always_inline)) inline void TensorDivRepeatM( const AscendC::LocalTensor& dst, const AscendC::LocalTensor& src, const AscendC::LocalTensor& src1, uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n) { - PIPE_BARRIER(V); - for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { - div_v(dst[vadd_idx * FLOAT_VECTOR_SIZE], - src[vadd_idx * FLOAT_VECTOR_SIZE], - src1, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - div_v(dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src1, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); + PlatformTensorDivRepeatM(dst, src, src1, sub_m, qk_n, qk_round_n); } -// Tensor 行归约求最大值（cmax_v + max_v 循环 + tail 处理） +// Tensor 行归约求最大值（委托 arch 层平台函数） __aicore__ __attribute__((always_inline)) inline void ReduceMaxRepeatM( const AscendC::LocalTensor& dst, const AscendC::LocalTensor& src, @@ -1922,72 +1212,10 @@ __aicore__ __attribute__((always_inline)) inline void ReduceMaxRepeatM( uint32_t qk_n, uint32_t qk_round_n) { - if (qk_n <= FLOAT_VECTOR_SIZE) { - __set_mask(qk_n); - cmax_v(dst, - src, - sub_m, // repeat - 1, // dstRepeatStride - 1, // srcBlockStride - qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride - ); - } else { - ub_to_ub( - tempTensor, - src, - 0, // sid - sub_m, // nBurst - HALF_VECTOR_SIZE / BLOCK_SIZE, // lenBurst - (qk_round_n - FLOAT_VECTOR_SIZE) / FLOAT_BLOCK_SIZE, // srcGap - 0 // dstGap - ); - PIPE_BARRIER(V); - for (uint32_t rowmax_idx = 1; rowmax_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowmax_idx) { - max_v( - tempTensor, - tempTensor, - src[rowmax_idx * FLOAT_VECTOR_SIZE], - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride - ); - PIPE_BARRIER(V); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - max_v( - tempTensor, - tempTensor, - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride - ); - } - PIPE_BARRIER(V); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - cmax_v( - dst, - tempTensor, - sub_m, // repeat - 1, // dstRepeatStride - 1, // srcBlockStride - 8 // srcRepeatStride - ); - } - SetVectorMask((uint64_t)-1, (uint64_t)-1); - PIPE_BARRIER(V); + PlatformReduceMaxRepeatM(dst, src, tempTensor, sub_m, qk_n, qk_round_n); } -// Tensor 行归约求和（cadd_v + add_v 循环 + tail 处理） +// Tensor 行归约求和（委托 arch 层平台函数） __aicore__ __attribute__((always_inline)) inline void ReduceSumRepeatM( const AscendC::LocalTensor& dst, const AscendC::LocalTensor& src, @@ -1995,131 +1223,34 @@ __aicore__ __attribute__((always_inline)) inline void ReduceSumRepeatM( uint32_t qk_n, uint32_t qk_round_n) { - if (qk_n <= FLOAT_VECTOR_SIZE) { - __set_mask(qk_n); - cadd_v( - dst, - src, - sub_m, // repeat - 1, // dstRepeatStride - 1, // srcBlockStride - qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } else { - for (uint32_t rowsum_idx = 1; rowsum_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowsum_idx) { - add_v( - src, - src, - src[rowsum_idx * FLOAT_VECTOR_SIZE], - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride - ); - PIPE_BARRIER(V); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - add_v( - src, - src, - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); - - cadd_v( - dst, - src, - sub_m, // repeat - 1, // dstRepeatStride - 1, // srcBlockStride - qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride - ); - } + PlatformReduceSumRepeatM(dst, src, sub_m, qk_n, qk_round_n); } -// Tensor 逐行乘法（mul_v 循环 + tail 处理） +// Tensor 逐行乘法（委托 arch 层平台函数） __aicore__ __attribute__((always_inline)) inline void TensorMulRepeatM( const AscendC::LocalTensor& dst, const AscendC::LocalTensor& src, const AscendC::LocalTensor& src1, uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t src1BlockStride ) { - PIPE_BARRIER(V); - for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { - mul_v(dst[vadd_idx * FLOAT_VECTOR_SIZE], - src[vadd_idx * FLOAT_VECTOR_SIZE], - src1, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - src1BlockStride, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - mul_v(dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src1, - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - src1BlockStride, // src1BlockStride - qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride - qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); + PlatformTensorMulRepeatM(dst, src, src1, sub_m, qk_n, qk_round_n, src1BlockStride); } // ====== 量化/反量化业务函数 ====== // ---- DeQuantPerHeadImpl 业务子函数 ---- -// 业务子函数1：加载 deScale + online 乘 quantScale -// 平台调用：PlatformMte2ToVSyncEvent2 / TensorMulRepeatM +// 业务子函数1：加载 deScale + online 乘 quantScale（委托arch层） __aicore__ __attribute__((always_inline)) inline void LoadDeScaleAndOnlineMul( const AscendC::GlobalTensor& deScaleGm, AscendC::LocalTensor deScaleUb, AscendC::LocalTensor quantScale, uint32_t sub_m, bool online) { - gm_to_ub_align(deScaleUb, - deScaleGm, - 0, // sid - 1, // nBurst - sub_m * sizeof(mmScaleType), // lenBurst - 0, // leftPaddingNum - 0, // rightPaddingNum - 0, // srcGap - 0 // dstGap - ); - if (online) { - PlatformMte2ToVSyncEvent2(); - TensorMulRepeatM(deScaleUb, deScaleUb, quantScale, 1, sub_m, RoundUp<16>(sub_m), 1); - } + PlatformLoadDeScaleAndOnlineMul(deScaleGm, deScaleUb, quantScale, sub_m, online); } -// 业务子函数2：加载 src(int32) + brcb 广播 deScale → tempScale -// 平台调用：PlatformMte2ToVSyncEvent0 / brcb_v / PlatformVPipeBarrier +// 业务子函数2：加载 src(int32) + brcb 广播 deScale → tempScale（委托arch层） __aicore__ __attribute__((always_inline)) inline void LoadSrcAndBrcbScale( const AscendC::GlobalTensor& src, AscendC::LocalTensor temp, @@ -2127,67 +1258,20 @@ __aicore__ __attribute__((always_inline)) inline void LoadSrcAndBrcbScale( AscendC::LocalTensor tempScale, uint32_t sub_m, uint32_t qk_round_n, bool move_tensor) { - if (move_tensor) { - gm_to_ub( - temp, - src, - 0, // sid - 1, // nBurst - CeilDiv(sub_m * qk_round_n), // lenBurst - 0, // srcGap - 0 // dstGap - ); - } - PlatformMte2ToVSyncEvent0(); - brcb_v( - tempScale.template ReinterpretCast(), - deScaleUb.template ReinterpretCast(), - 1, // dstBlockStrides - 8, // dstRepeatStride - RoundUp<16>(sub_m) / FLOAT_BLOCK_SIZE // repeat - ); - PlatformVPipeBarrier(); + PlatformLoadSrcAndBrcbScale(src, temp, deScaleUb, tempScale, sub_m, qk_round_n, move_tensor); } -// 业务子函数3：INT32→FP32 转换(conv_v) + 乘 tempScale -// 平台调用：conv_v / TensorMulRepeatM / PlatformVPipeBarrier +// 业务子函数3：INT32→FP32 转换(conv_v) + 乘 tempScale（委托arch层） __aicore__ __attribute__((always_inline)) inline void ConvInt32ToFP32AndMul( AscendC::LocalTensor dst, AscendC::LocalTensor temp, AscendC::LocalTensor tempScale, uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n) { - uint32_t count = sub_m * qk_round_n; - uint32_t repeat_times = (count + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE; - if (repeat_times < 255) { - conv_v( - dst, // dst - temp, // src - repeat_times, // repeat_times - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - } else { - for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { - conv_v( - dst[vconv_idx * count / 2], // dst - temp[vconv_idx * count / 2], // src - (count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat_times - 1, // dstBlockStride - 1, // srcBlockStride - 8, // dstRepeatStride - 8 // srcRepeatStride - ); - } - } - TensorMulRepeatM(dst, dst, tempScale, sub_m, qk_n, qk_round_n, 0); - PlatformVPipeBarrier(); + PlatformConvInt32ToFP32AndMul(dst, temp, tempScale, sub_m, qk_n, qk_round_n); } -// 逐 Head 反量化：INT32 → FP32 + deScale 乘法（编排层） -// 依赖：LoadDeScaleAndOnlineMul / LoadSrcAndBrcbScale / ConvInt32ToFP32AndMul +// 逐 Head 反量化：INT32 → FP32 + deScale 乘法（委托arch层编排） __aicore__ __attribute__((always_inline)) inline void DeQuantPerHeadImpl( const AscendC::GlobalTensor& deScaleGm, const AscendC::GlobalTensor& src, @@ -2202,64 +1286,16 @@ __aicore__ __attribute__((always_inline)) inline void DeQuantPerHeadImpl( bool online, bool move_tensor ){ - LoadDeScaleAndOnlineMul(deScaleGm, deScaleUb, quantScale, sub_m, online); - LoadSrcAndBrcbScale(src, temp, deScaleUb, tempScale, sub_m, qk_round_n, move_tensor); - ConvInt32ToFP32AndMul(dst, temp, tempScale, sub_m, qk_n, qk_round_n); + PlatformDeQuantPerHead(deScaleGm, src, dst, temp, deScaleUb, tempScale, quantScale, + sub_m, qk_n, qk_round_n, online, move_tensor); } -// 逐 Token 量化：FP32 → FP16 → INT8 -// 依赖：TensorDivRepeatM / TensorMulRepeatM / conv_v / Cast +// 逐 Token 量化：FP32 → FP16 → INT8（委托arch层） __aicore__ __attribute__((always_inline)) inline void QuantPerTokenImpl( const AscendC::LocalTensor& dst, const AscendC::LocalTensor& src, const AscendC::LocalTensor& scale, uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t pQuantOnline) { - if (pQuantOnline) { - // scr / scale - TensorDivRepeatM(dst.template ReinterpretCast(), src, scale, sub_m, qk_n, qk_round_n); - } else { - // scr * scale - TensorMulRepeatM(dst.template ReinterpretCast(), src, scale, sub_m, qk_n, qk_round_n, 0); - } - // src fp32 -> casttofp16 -> casttoint8 - uint32_t count = sub_m * qk_round_n; - uint32_t repeat_times = (count + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE; - if (repeat_times < 255) { - conv_v( - dst.template ReinterpretCast(), // dst - dst.template ReinterpretCast(), // src - repeat_times, // repeat_times - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); - } else { - for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { - conv_v( - dst.template ReinterpretCast()[vconv_idx * count / 2], // dst - dst.template ReinterpretCast()[vconv_idx * count / 2], // src - (count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat_times - 1, // dstBlockStride - 1, // srcBlockStride - 4, // dstRepeatStride - 8 // srcRepeatStride - ); - } - } - PIPE_BARRIER(V); - for (uint32_t row_idx = 0; row_idx < qk_n / HALF_VECTOR_SIZE; ++row_idx) { - AscendC::Cast(dst.template ReinterpretCast()[row_idx * HALF_VECTOR_SIZE], - dst.template ReinterpretCast()[row_idx * HALF_VECTOR_SIZE], AscendC::RoundMode::CAST_RINT, - (uint64_t)0, sub_m, {1, 1, (uint8_t)((qk_round_n) / BLOCK_SIZE), (uint8_t)(qk_round_n / BLOCK_SIZE)}); - } - if (qk_n % HALF_VECTOR_SIZE > 0) { - __set_mask(qk_n % HALF_VECTOR_SIZE); - AscendC::Cast(dst.template ReinterpretCast()[qk_n / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], - dst.template ReinterpretCast()[qk_n / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], AscendC::RoundMode::CAST_RINT, - (uint64_t)0, sub_m, {1, 1, (uint8_t)((qk_round_n) / BLOCK_SIZE), (uint8_t)(qk_round_n / BLOCK_SIZE)}); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); + PlatformQuantPerToken(dst, src, scale, sub_m, qk_n, qk_round_n, pQuantOnline); } \ No newline at end of file From 14614a3b6b4fb28e57d44149de92521629fb3da0 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Thu, 13 Aug 2026 19:02:17 +0800 Subject: [PATCH 18/24] feat: youhua jiagou a3 --- .../op_kernel/multi_latent_attention.h | 243 +------ .../multi_latent_attention_aiv_arch32.h | 132 ++-- .../op_kernel/multi_latent_attention_aiv_bs.h | 10 +- .../op_kernel/multi_latent_attention_bs.h | 621 ------------------ 4 files changed, 77 insertions(+), 929 deletions(-) diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h index 6a5465a..af4d63b 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h @@ -1387,250 +1387,15 @@ class MLADecoderAiv{ private: // ====== AIV refactor: context struct + sub-functions ====== - - struct VectorContext { - // input params - uint32_t cur_batch; - uint32_t start_head; - uint32_t cur_nIndx; - uint32_t cur_q_seqlen; - uint32_t cur_kv_seqlen; - uint32_t cur_head_num; - uint32_t offset_tiling; - - // addresses - uint64_t addr_o_scalar; - uint64_t addr_mask_scalar; - uint32_t mask_offset; - - // loop & size - uint32_t pp_n_scalar; - uint32_t sub_n_loop; - uint32_t real_n_loop; - uint32_t n_loop; - - // QK dims - uint32_t qk_n; - uint32_t qk_round_n; - uint32_t qk_n_2; - uint32_t qk_round_n_2; - - // head split - uint32_t sub_head_num; - uint32_t sub_m; - uint32_t head_idx; - uint64_t o_offset; - - // tail info - uint32_t tail_len; - bool prev_tail_mask; - }; - - __aicore__ __attribute__((always_inline)) inline void InitVectorContext( - VectorContext &ctx, uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, - uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, - uint32_t offset_tiling) - { - ctx.cur_batch = cur_batch; - ctx.start_head = start_head; - ctx.cur_nIndx = cur_nIndx; - ctx.cur_q_seqlen = cur_q_seqlen; - ctx.cur_kv_seqlen = cur_kv_seqlen; - ctx.cur_head_num = cur_head_num; - ctx.offset_tiling = offset_tiling; - - // addr - uint32_t addr_o_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 4 + offset_tiling)); - uint32_t addr_o_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 5 + offset_tiling)); - ctx.addr_o_scalar = (uint64_t)(((uint64_t)addr_o_high32) << 32 | addr_o_loww32); - - uint32_t addr_mask_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 6 + offset_tiling)); - uint32_t addr_mask_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 7 + offset_tiling)); - ctx.addr_mask_scalar = (uint64_t)(((uint64_t)addr_mask_high32) << 32 | addr_mask_loww32); - ctx.mask_offset = ctx.addr_mask_scalar; - - // loop & size - ctx.pp_n_scalar = block_size; - ctx.sub_n_loop = ctx.pp_n_scalar / block_size; - ctx.real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; - ctx.n_loop = (cur_kv_seqlen + ctx.pp_n_scalar - 1) / ctx.pp_n_scalar; - - // QK dims - ctx.qk_n = ctx.pp_n_scalar; - ctx.qk_round_n = RoundUp(ctx.qk_n); - ctx.qk_n_2 = ctx.pp_n_scalar; - ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); - - // head split - ctx.sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; - ctx.sub_m = ctx.sub_head_num * cur_q_seqlen; - ctx.head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; - ctx.o_offset = ctx.addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; - - // tail info - ctx.tail_len = cur_kv_seqlen - (ctx.n_loop - 1) * ctx.pp_n_scalar; - ctx.prev_tail_mask = (ctx.n_loop > 1 && ctx.tail_len < cur_q_seqlen - 1); - } - - struct VectorTP1Context { - // input params - uint32_t cur_batch; - uint32_t start_head; - uint32_t cur_nIndx; - uint32_t cur_q_seqlen; - uint32_t cur_kv_seqlen; - uint32_t cur_head_num; - uint32_t offset_tiling; - - // addresses - uint64_t addr_o_scalar; - uint64_t addr_mask_scalar; - uint32_t mask_offset; - - // loop & size - uint32_t pp_n_scalar; - uint32_t sub_n_loop; - uint32_t real_n_loop; - uint32_t n_loop; - - // QK dims - uint32_t qk_n; - uint32_t qk_round_n; - uint32_t qk_n_2; - uint32_t qk_round_n_2; - - // head split - uint32_t sub_head_num; - uint32_t sub_m; - uint32_t head_idx; - uint64_t o_offset; - - // TP1 specific - uint32_t s_block_stack; - uint32_t m_slice; - uint32_t m_end; - }; - - __aicore__ __attribute__((always_inline)) inline void InitVectorTP1Context( - VectorTP1Context &ctx, uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, - uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, - uint32_t offset_tiling) - { - ctx.cur_batch = cur_batch; - ctx.start_head = start_head; - ctx.cur_nIndx = cur_nIndx; - ctx.cur_q_seqlen = cur_q_seqlen; - ctx.cur_kv_seqlen = cur_kv_seqlen; - ctx.cur_head_num = cur_head_num; - ctx.offset_tiling = offset_tiling; - - // addr - uint32_t prev_task = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); - ctx.addr_o_scalar = prev_task * q_heads * embedding_size; - ctx.addr_mask_scalar = 0; - ctx.mask_offset = 0; - - // loop & size - ctx.pp_n_scalar = block_size; - ctx.sub_n_loop = ctx.pp_n_scalar / block_size; - ctx.real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; - ctx.n_loop = (cur_kv_seqlen + ctx.pp_n_scalar - 1) / ctx.pp_n_scalar; - - // QK dims - ctx.qk_n = ctx.pp_n_scalar; - ctx.qk_round_n = RoundUp(ctx.qk_n); - ctx.qk_n_2 = ctx.pp_n_scalar; - ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); - - // head split - ctx.sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; - ctx.sub_m = ctx.sub_head_num * cur_q_seqlen; - ctx.head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; - ctx.o_offset = ctx.addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; - - // TP1 specific - ctx.s_block_stack = 4; - ctx.m_slice = FLOAT_VECTOR_SIZE / ctx.s_block_stack; - ctx.m_end = (ctx.sub_m + ctx.m_slice - 1) / ctx.m_slice; - } + // VectorContext / InitVectorContext / VectorTP1Context / InitVectorTP1Context + // 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） // [AIV 第三层平台函数 + 第二层业务函数] 已迁移至 aiv_arch32.h / aiv_bs.h #include "multi_latent_attention_aiv_arch32.h" #include "multi_latent_attention_aiv_bs.h" - - __aicore__ __attribute__((always_inline)) inline void ReduceMaxRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& tempTensor, - uint32_t sub_m, - uint32_t qk_n, - uint32_t qk_round_n) - { - if (qk_n <= FLOAT_VECTOR_SIZE) { - __set_mask(qk_n); - cmax_v(dst, - src, - sub_m, // repeat - 1, // dstRepeatStride - 1, // srcBlockStride - qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride - ); - } else { - ub_to_ub( - tempTensor, - src, - 0, // sid - sub_m, // nBurst - HALF_VECTOR_SIZE / BLOCK_SIZE, // lenBurst - (qk_round_n - FLOAT_VECTOR_SIZE) / FLOAT_BLOCK_SIZE, // srcGap - 0 // dstGap - ); - PIPE_BARRIER(V); - for (uint32_t rowmax_idx = 1; rowmax_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowmax_idx) { - max_v( - tempTensor, - tempTensor, - src[rowmax_idx * FLOAT_VECTOR_SIZE], - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride - ); - PIPE_BARRIER(V); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - max_v( - tempTensor, - tempTensor, - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 1, // src1BlockStride - 8, // dstRepeatStride - 8, // src0RepeatStride - qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride - ); - } - PIPE_BARRIER(V); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - cmax_v( - dst, - tempTensor, - sub_m, // repeat - 1, // dstRepeatStride - 1, // srcBlockStride - 8 // srcRepeatStride - ); - } - SetVectorMask((uint64_t)-1, (uint64_t)-1); - PIPE_BARRIER(V); - } + // ReduceMaxRepeatM + // 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） // ReduceSumRepeatM diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h index d1fbf15..8477643 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h @@ -10,6 +10,8 @@ // 本文件仅被 MLADecoderAiv 类 include，包含 Vector 侧平台同步函数。 // AIC 侧平台函数在 multi_latent_attention_arch32.h 中。 +#pragma once + // 平台函数：非TP1 SoftmaxStage1 前置管道同步 // 在 SoftmaxStage1 调用之前执行：等待 QK 就绪 + 等待 MTE3→MTE2 管道同步 __aicore__ __attribute__((always_inline)) inline void PlatformSoftmaxStage1PreSync() @@ -134,7 +136,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformMulVectorByBroadca dm32_ubuf_tensor.ReinterpretCast(), 1, // dstBlockStride 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat + (uint8_t)(round_sub_m / FLOAT_BLOCK_SIZE) // repeat ); PIPE_BARRIER(V); @@ -143,13 +145,13 @@ __aicore__ __attribute__((always_inline)) inline void PlatformMulVectorByBroadca for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat + tv32_ubuf_tensor.template ReinterpretCast(), + (uint8_t)sub_m, // repeat 1, // dstBlockStride 1, // src0BlockStride 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // dstRepeatStride + (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // src0RepeatStride 1 // src1RepeatStride ); } @@ -158,13 +160,13 @@ __aicore__ __attribute__((always_inline)) inline void PlatformMulVectorByBroadca __set_mask(__v % FLOAT_VECTOR_SIZE); mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat + tv32_ubuf_tensor.template ReinterpretCast(), + (uint8_t)sub_m, // repeat 1, // dstBlockStride 1, // src0BlockStride 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // dstRepeatStride + (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // src0RepeatStride 1 // src1RepeatStride ); SetVectorMask((uint64_t)-1, (uint64_t)-1); @@ -190,7 +192,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformDivVectorByBroadca gl32_ubuf_tensor.ReinterpretCast()[head_loop_idx * 16], 1, // dstBlockStride 8, // dstRepeatStride - round_sub_m / FLOAT_BLOCK_SIZE // repeat + (uint8_t)(round_sub_m / FLOAT_BLOCK_SIZE) // repeat ); PIPE_BARRIER(V); @@ -199,13 +201,13 @@ __aicore__ __attribute__((always_inline)) inline void PlatformDivVectorByBroadca for (uint32_t vdiv_idx = 0; vdiv_idx < __v / FLOAT_VECTOR_SIZE; ++vdiv_idx) { div_v(go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat + tv32_ubuf_tensor.template ReinterpretCast(), + (uint8_t)sub_m, // repeat 1, // dstBlockStride 1, // src0BlockStride 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // dstRepeatStride + (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // src0RepeatStride 1 // src1RepeatStride ); } @@ -214,13 +216,13 @@ __aicore__ __attribute__((always_inline)) inline void PlatformDivVectorByBroadca __set_mask(__v % FLOAT_VECTOR_SIZE); div_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor, - sub_m, // repeat + tv32_ubuf_tensor.template ReinterpretCast(), + (uint8_t)sub_m, // repeat 1, // dstBlockStride 1, // src0BlockStride 0, // src1BlockStride - round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride - round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // dstRepeatStride + (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // src0RepeatStride 1 // src1RepeatStride ); SetVectorMask((uint64_t)-1, (uint64_t)-1); // fix hidden_size=96 @@ -661,7 +663,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformMaxV( uint32_t repeat) { max_v( - dst, src0, src1, repeat, 1, 1, 1, 8, 8, 8); + dst, src0, src1, (uint8_t)repeat, 1, 1, 1, 8, 8, 8); } // 平台函数：sub_v（逐元素减法） @@ -672,7 +674,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformSubV( uint32_t repeat) { sub_v( - dst, src0, src1, repeat, 1, 1, 1, 8, 8, 8); + dst, src0, src1, (uint8_t)repeat, 1, 1, 1, 8, 8, 8); } // 平台函数：exp_v（指数运算） @@ -682,7 +684,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformExpV( uint32_t repeat) { exp_v( - dst, src, repeat, 1, 1, 8, 8); + dst, src, (uint8_t)repeat, 1, 1, 8, 8); } // 平台函数：mul_v（逐元素乘法） @@ -693,7 +695,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformMulV( uint32_t repeat) { mul_v( - dst, src0, src1, repeat, 1, 1, 1, 8, 8, 8); + dst, src0, src1, (uint8_t)repeat, 1, 1, 1, 8, 8, 8); } // 平台函数：add_v（逐元素加法） @@ -704,7 +706,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformAddV( uint32_t repeat) { add_v( - dst, src0, src1, repeat, 1, 1, 1, 8, 8, 8); + dst, src0, src1, (uint8_t)repeat, 1, 1, 1, 8, 8, 8); } // 平台函数：muls_v（标量乘法，单次调用） @@ -715,7 +717,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformMulsV( uint32_t repeat) { muls_v( - dst, src, scalar, repeat, 1, 1, 8, 8); + dst, src, scalar, (uint8_t)repeat, 1, 1, 8, 8); } // 平台函数：brcb_v（广播，uint32_t） @@ -724,7 +726,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformBrcbV( AscendC::LocalTensor src, uint32_t repeat) { - brcb_v(dst, src, 1, 8, repeat); + brcb_v(dst, src, 1, 8, (uint8_t)repeat); } // 平台函数：brcb_v（广播，float） @@ -736,7 +738,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformBrcbVFloat( brcb_v( dst.ReinterpretCast(), src.ReinterpretCast(), - 1, 8, round_sub_m / FLOAT_BLOCK_SIZE); + 1, 8, (uint8_t)(round_sub_m / FLOAT_BLOCK_SIZE)); } // 平台函数：conv_v（float→OUT_DTYPE 转换） @@ -746,7 +748,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformConvVToOutput( uint32_t repeat) { conv_v( - dst, src, repeat, 1, 1, 4, 8); + dst, src, (uint8_t)repeat, 1, 1, 4, 8); } // 平台函数：conv_v（int32_t→float 转换，含 repeat_times<255 分支） @@ -758,12 +760,12 @@ __aicore__ __attribute__((always_inline)) inline void PlatformConvInt32ToFloat( uint32_t repeat_times = (count + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE; if (repeat_times < 255) { conv_v( - dst, temp, repeat_times, 1, 1, 8, 8); + dst, temp, (uint8_t)repeat_times, 1, 1, 8, 8); } else { for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { conv_v( dst[vconv_idx * count / 2], temp[vconv_idx * count / 2], - (count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + (uint8_t)((count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE), 1, 1, 8, 8); } } @@ -779,13 +781,13 @@ __aicore__ __attribute__((always_inline)) inline void PlatformConvFloatToHalf( conv_v( dst.template ReinterpretCast(), dst.template ReinterpretCast(), - repeat_times, 1, 1, 4, 8); + (uint8_t)repeat_times, 1, 1, 4, 8); } else { for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { conv_v( dst.template ReinterpretCast()[vconv_idx * count / 2], dst.template ReinterpretCast()[vconv_idx * count / 2], - (count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + (uint8_t)((count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE), 1, 1, 4, 8); } } @@ -837,24 +839,24 @@ __aicore__ __attribute__((always_inline)) inline void PlatformTensorSubValueRepe brcb_v( tempMaxTensor.ReinterpretCast(), MaxTensor.ReinterpretCast(), - 1, 8, round_sub_m / FLOAT_BLOCK_SIZE); + 1, 8, (uint8_t)(round_sub_m / FLOAT_BLOCK_SIZE)); PIPE_BARRIER(V); for (uint32_t sub_v_idx = 0; sub_v_idx < qk_n / FLOAT_VECTOR_SIZE; ++sub_v_idx) { sub_v( dst[sub_v_idx * FLOAT_VECTOR_SIZE], src[sub_v_idx * FLOAT_VECTOR_SIZE], - tempMaxTensor, sub_m, 1, 1, 0, - qk_round_n / FLOAT_BLOCK_SIZE, - qk_round_n / FLOAT_BLOCK_SIZE, 1); + tempMaxTensor, (uint8_t)sub_m, 1, 1, 0, + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), 1); } if (qk_n % FLOAT_VECTOR_SIZE > 0) { __set_mask(qk_n % FLOAT_VECTOR_SIZE); sub_v( dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tempMaxTensor, sub_m, 1, 1, 0, - qk_round_n / FLOAT_BLOCK_SIZE, - qk_round_n / FLOAT_BLOCK_SIZE, 1); + tempMaxTensor, (uint8_t)sub_m, 1, 1, 0, + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), 1); SetVectorMask((uint64_t)-1, (uint64_t)-1); } PIPE_BARRIER(V); @@ -872,18 +874,18 @@ __aicore__ __attribute__((always_inline)) inline void PlatformTensorDivRepeatM( div_v( dst[vadd_idx * FLOAT_VECTOR_SIZE], src[vadd_idx * FLOAT_VECTOR_SIZE], - src1, sub_m, 1, 1, 0, - qk_round_n / FLOAT_BLOCK_SIZE, - qk_round_n / FLOAT_BLOCK_SIZE, 1); + src1, (uint8_t)sub_m, 1, 1, 0, + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), 1); } if (qk_n % FLOAT_VECTOR_SIZE > 0) { __set_mask(qk_n % FLOAT_VECTOR_SIZE); div_v( dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src1, sub_m, 1, 1, 0, - qk_round_n / FLOAT_BLOCK_SIZE, - qk_round_n / FLOAT_BLOCK_SIZE, 1); + src1, (uint8_t)sub_m, 1, 1, 0, + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), 1); SetVectorMask((uint64_t)-1, (uint64_t)-1); } PIPE_BARRIER(V); @@ -901,7 +903,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformReduceMaxRepeatM( if (qk_n <= FLOAT_VECTOR_SIZE) { __set_mask(qk_n); cmax_v( - dst, src, sub_m, 1, 1, qk_round_n / FLOAT_BLOCK_SIZE); + dst, src, (uint8_t)sub_m, 1, 1, (uint16_t)(qk_round_n / FLOAT_BLOCK_SIZE)); } else { ub_to_ub( tempTensor, src, 0, sub_m, @@ -911,7 +913,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformReduceMaxRepeatM( for (uint32_t rowmax_idx = 1; rowmax_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowmax_idx) { max_v( tempTensor, tempTensor, src[rowmax_idx * FLOAT_VECTOR_SIZE], - sub_m, 1, 1, 1, 8, 8, qk_round_n / FLOAT_BLOCK_SIZE); + (uint8_t)sub_m, 1, 1, 1, 8, 8, (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE)); PIPE_BARRIER(V); } if (qk_n % FLOAT_VECTOR_SIZE > 0) { @@ -919,12 +921,12 @@ __aicore__ __attribute__((always_inline)) inline void PlatformReduceMaxRepeatM( max_v( tempTensor, tempTensor, src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - sub_m, 1, 1, 1, 8, 8, qk_round_n / FLOAT_BLOCK_SIZE); + (uint8_t)sub_m, 1, 1, 1, 8, 8, (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE)); } PIPE_BARRIER(V); SetVectorMask((uint64_t)-1, (uint64_t)-1); cmax_v( - dst, tempTensor, sub_m, 1, 1, 8); + dst, tempTensor, (uint8_t)sub_m, 1, 1, 8); } SetVectorMask((uint64_t)-1, (uint64_t)-1); PIPE_BARRIER(V); @@ -941,16 +943,16 @@ __aicore__ __attribute__((always_inline)) inline void PlatformReduceSumRepeatM( if (qk_n <= FLOAT_VECTOR_SIZE) { __set_mask(qk_n); cadd_v( - dst, src, sub_m, 1, 1, qk_round_n / FLOAT_BLOCK_SIZE); + dst, src, (uint8_t)sub_m, 1, 1, (uint16_t)(qk_round_n / FLOAT_BLOCK_SIZE)); SetVectorMask((uint64_t)-1, (uint64_t)-1); } else { for (uint32_t rowsum_idx = 1; rowsum_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowsum_idx) { add_v( src, src, src[rowsum_idx * FLOAT_VECTOR_SIZE], - sub_m, 1, 1, 1, - qk_round_n / FLOAT_BLOCK_SIZE, - qk_round_n / FLOAT_BLOCK_SIZE, - qk_round_n / FLOAT_BLOCK_SIZE); + (uint8_t)sub_m, 1, 1, 1, + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE)); PIPE_BARRIER(V); } if (qk_n % FLOAT_VECTOR_SIZE > 0) { @@ -958,15 +960,15 @@ __aicore__ __attribute__((always_inline)) inline void PlatformReduceSumRepeatM( add_v( src, src, src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - sub_m, 1, 1, 1, - qk_round_n / FLOAT_BLOCK_SIZE, - qk_round_n / FLOAT_BLOCK_SIZE, - qk_round_n / FLOAT_BLOCK_SIZE); + (uint8_t)sub_m, 1, 1, 1, + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE)); SetVectorMask((uint64_t)-1, (uint64_t)-1); } PIPE_BARRIER(V); cadd_v( - dst, src, sub_m, 1, 1, qk_round_n / FLOAT_BLOCK_SIZE); + dst, src, (uint8_t)sub_m, 1, 1, (uint16_t)(qk_round_n / FLOAT_BLOCK_SIZE)); } } @@ -982,18 +984,18 @@ __aicore__ __attribute__((always_inline)) inline void PlatformTensorMulRepeatM( mul_v( dst[vadd_idx * FLOAT_VECTOR_SIZE], src[vadd_idx * FLOAT_VECTOR_SIZE], - src1, sub_m, 1, 1, src1BlockStride, - qk_round_n / FLOAT_BLOCK_SIZE, - qk_round_n / FLOAT_BLOCK_SIZE, 1); + src1, (uint8_t)sub_m, 1, 1, (uint8_t)src1BlockStride, + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), 1); } if (qk_n % FLOAT_VECTOR_SIZE > 0) { __set_mask(qk_n % FLOAT_VECTOR_SIZE); mul_v( dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src1, sub_m, 1, 1, src1BlockStride, - qk_round_n / FLOAT_BLOCK_SIZE, - qk_round_n / FLOAT_BLOCK_SIZE, 1); + src1, (uint8_t)sub_m, 1, 1, (uint8_t)src1BlockStride, + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), + (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), 1); SetVectorMask((uint64_t)-1, (uint64_t)-1); } PIPE_BARRIER(V); @@ -1032,7 +1034,7 @@ __aicore__ __attribute__((always_inline)) inline void PlatformLoadSrcAndBrcbScal brcb_v( tempScale.template ReinterpretCast(), deScaleUb.template ReinterpretCast(), - 1, 8, RoundUp<16>(sub_m) / FLOAT_BLOCK_SIZE); + 1, 8, (uint8_t)(RoundUp<16>(sub_m) / FLOAT_BLOCK_SIZE)); PlatformVPipeBarrier(); } diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h index 2c77ad7..2240fc5 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h @@ -10,6 +10,8 @@ // 本文件仅被 MLADecoderAiv 类 include，包含 Vector 侧业务调度函数。 // AIC 侧业务函数在 multi_latent_attention_bs.h 中。 +#pragma once + // ====== AIV refactor: context struct + init functions ====== // 非 TP1 路径上下文 @@ -497,7 +499,7 @@ __aicore__ __attribute__((always_inline)) inline void Stage2MergeAccumulate( AscendC::GlobalTensor go_gm_tensor, AscendC::LocalTensor go32_ubuf_tensor, AscendC::LocalTensor lo_ubuf_tensor, - AscendC::LocalTensor tv32_ubuf_tensor, + AscendC::LocalTensor tv32_ubuf_tensor, AscendC::LocalTensor dm32_ubuf_tensor, uint32_t sub_m, uint32_t round_sub_m, @@ -505,7 +507,7 @@ __aicore__ __attribute__((always_inline)) inline void Stage2MergeAccumulate( uint32_t head_loop) { // brcb dm → tv, go = go * dm_block - PlatformMulVectorByBroadcast(go32_ubuf_tensor, tv32_ubuf_tensor, + PlatformMulVectorByBroadcast(go32_ubuf_tensor, tv32_ubuf_tensor.template ReinterpretCast(), dm32_ubuf_tensor, sub_m, round_sub_m, round_v, __v); if (head_loop > 1) { @@ -534,7 +536,7 @@ __aicore__ __attribute__((always_inline)) inline void Stage2FinalizeAndOutput( AscendC::GlobalTensor go_gm_tensor, AscendC::LocalTensor go32_ubuf_tensor, AscendC::LocalTensor go_ubuf_tensor, - AscendC::LocalTensor tv32_ubuf_tensor, + AscendC::LocalTensor tv32_ubuf_tensor, AscendC::LocalTensor gl32_ubuf_tensor, AscendC::LocalTensor lse32_ubuf_tensor, AscendC::LocalTensor lse_conv_ubuf_tensor, @@ -555,7 +557,7 @@ __aicore__ __attribute__((always_inline)) inline void Stage2FinalizeAndOutput( uint32_t sub_m_d64 = (sub_m + 63) / 64; // go = go / gl_block - PlatformDivVectorByBroadcast(go32_ubuf_tensor, tv32_ubuf_tensor, + PlatformDivVectorByBroadcast(go32_ubuf_tensor, tv32_ubuf_tensor.template ReinterpretCast(), gl32_ubuf_tensor, sub_m, round_sub_m, round_v, __v, head_loop_idx); // go = castfp32to16(go) + V→MTE3 同步 diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h index 8b91ac4..7ca42de 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h @@ -247,130 +247,6 @@ __aicore__ __attribute__((always_inline)) inline void LoadQRopeFromGMToL1( } } -// ==================== TP1 QK 业务函数 ==================== - -// TP1 QK 参数结构体 -struct TP1QKParams { - uint32_t qk_n; - uint32_t qk_round_n; - uint32_t l1_kv_pingpong_flag; - uint32_t embed_split_size; - uint32_t round_embed_split_size; - int64_t now_l1_offset; - int64_t kv_offset; - int64_t kv_offset_rope; - uint32_t sv_round_n; - uint32_t gm_split_idx; // split_idx 用于 GM 偏移计算 - uint32_t pp_n_scalar; // block_size，用于 GM 偏移计算 -}; - -// 业务函数：初始化 TP1 QK 参数 -// 对应原始 CUBE1 stage1 中 split_idx 循环体内的参数计算 -__aicore__ __attribute__((always_inline)) inline void InitTP1QKParams( - TP1Context &ctx, uint32_t split_idx, uint32_t n_idx, TP1QKParams ¶ms) -{ - uint32_t pp_n_scalar = ctx.pp_n_scalar; - uint32_t n_loop = ctx.n_loop; - uint32_t now_idx = n_idx + split_idx; - - params.l1_kv_pingpong_flag = now_idx % 2; - - // 动态计算 qk_n / qk_round_n - if (now_idx == (n_loop - 1)) { - params.qk_n = (ctx.cur_kv_seqlen - now_idx * pp_n_scalar); - } else { - params.qk_n = pp_n_scalar; - } - params.qk_round_n = RoundUp(params.qk_n); - - // embed_split_size 默认 128，idx==4 时为 64（在编排层设置） - params.embed_split_size = 128; - params.round_embed_split_size = RoundUp(params.embed_split_size); - - // now_l1_offset 默认 0，在 LoadTP1KVMain/KVRope 中更新 - params.now_l1_offset = 0; - - // block_table_id / kv_offset / kv_offset_rope 地址计算 - uint32_t block_table_id = (uint32_t)(*(block_tables_gm + - ctx.cur_batch * max_num_blocks_per_query + ctx.start_kv / block_size + now_idx)); - params.kv_offset = (int64_t)block_table_id * block_size * stride_kv; - params.kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; - - // sv_round_n 计算（用于 S→GM 的 dstStride） - uint32_t sv_n_triu = n_loop * pp_n_scalar; - uint32_t sv_n; - if (n_idx + ctx.s_block_stack > n_loop - 1) { - sv_n = ctx.cur_kv_seqlen - n_idx * pp_n_scalar; - } else { - sv_n = pp_n_scalar * ctx.s_block_stack; - } - params.sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; - - params.gm_split_idx = split_idx; - params.pp_n_scalar = pp_n_scalar; -} - -// 业务函数：TP1 Q 数据从 L1 加载到 L0A -// 复用非 TP1 的 PlatformWaitQLoadComplete / PlatformLoadQToL0A / PlatformSetQLoadComplete -// TP1 中 q_load_coeff = m，与非 TP1 一致 -__aicore__ __attribute__((always_inline)) inline void LoadTP1QDataToL0A( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformWaitQLoadComplete(embed_split_idx); - PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, m, false); - PlatformSetQLoadComplete(embed_split_idx); -} - -// 业务函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2） -// 更新 now_l1_offset 供后续 L1→L0B 使用 -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVMainFromGM( - TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 256; - PlatformLoadTP1KVMainToL1(embed_split_idx, params.qk_n, params.qk_round_n, - params.l1_kv_pingpong_flag, params.kv_offset, - embed_split_idx * 128); -} - -// 业务函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4） -// 更新 now_l1_offset 供后续 L1→L0B 使用 -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVRopeFromGM( - TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; - PlatformLoadTP1KVRopeToL1(params.qk_n, params.qk_round_n, - params.l1_kv_pingpong_flag, params.kv_offset_rope); -} - -// 业务函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVDataToL0B( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformLoadTP1KVToL0B(embed_split_idx, params.round_embed_split_size, - params.qk_round_n, params.now_l1_offset, - params.l1_kv_pingpong_flag); -} - -// 业务函数：TP1 QK MMAD 计算 -__aicore__ __attribute__((always_inline)) inline void ComputeTP1QKMMad( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformComputeTP1QKMMad(embed_split_idx, params.embed_split_size, - m, params.qk_n, params.l1_kv_pingpong_flag); -} - -// 业务函数：TP1 QK 结果拷贝到 GM（idx 4 时） -// GM 偏移: block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar -__aicore__ __attribute__((always_inline)) inline void CopyTP1QKResultToGM( - const TP1QKParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) -{ - uint64_t gm_dst_offset = (uint64_t)block_idx * TMP_SIZE_DECODER * 4 + - (uint64_t)((n_idx / 4) % 2) * TMP_SIZE_DECODER * 2 + - params.gm_split_idx * params.pp_n_scalar; - PlatformCopyTP1QKResultToGM(m, params.qk_round_n, - params.l1_kv_pingpong_flag, - gm_dst_offset, params.sv_round_n); -} // 业务函数：将 KV 主体数据从 GM 搬运到 L1 // ND_FORMAT 场景：ND→NZ 格式转换（gm_to_l1） @@ -410,130 +286,6 @@ __aicore__ __attribute__((always_inline)) inline void LoadKVMainFromGMToL1( } } -// ==================== TP1 QK 业务函数 ==================== - -// TP1 QK 参数结构体 -struct TP1QKParams { - uint32_t qk_n; - uint32_t qk_round_n; - uint32_t l1_kv_pingpong_flag; - uint32_t embed_split_size; - uint32_t round_embed_split_size; - int64_t now_l1_offset; - int64_t kv_offset; - int64_t kv_offset_rope; - uint32_t sv_round_n; - uint32_t gm_split_idx; // split_idx 用于 GM 偏移计算 - uint32_t pp_n_scalar; // block_size，用于 GM 偏移计算 -}; - -// 业务函数：初始化 TP1 QK 参数 -// 对应原始 CUBE1 stage1 中 split_idx 循环体内的参数计算 -__aicore__ __attribute__((always_inline)) inline void InitTP1QKParams( - TP1Context &ctx, uint32_t split_idx, uint32_t n_idx, TP1QKParams ¶ms) -{ - uint32_t pp_n_scalar = ctx.pp_n_scalar; - uint32_t n_loop = ctx.n_loop; - uint32_t now_idx = n_idx + split_idx; - - params.l1_kv_pingpong_flag = now_idx % 2; - - // 动态计算 qk_n / qk_round_n - if (now_idx == (n_loop - 1)) { - params.qk_n = (ctx.cur_kv_seqlen - now_idx * pp_n_scalar); - } else { - params.qk_n = pp_n_scalar; - } - params.qk_round_n = RoundUp(params.qk_n); - - // embed_split_size 默认 128，idx==4 时为 64（在编排层设置） - params.embed_split_size = 128; - params.round_embed_split_size = RoundUp(params.embed_split_size); - - // now_l1_offset 默认 0，在 LoadTP1KVMain/KVRope 中更新 - params.now_l1_offset = 0; - - // block_table_id / kv_offset / kv_offset_rope 地址计算 - uint32_t block_table_id = (uint32_t)(*(block_tables_gm + - ctx.cur_batch * max_num_blocks_per_query + ctx.start_kv / block_size + now_idx)); - params.kv_offset = (int64_t)block_table_id * block_size * stride_kv; - params.kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; - - // sv_round_n 计算（用于 S→GM 的 dstStride） - uint32_t sv_n_triu = n_loop * pp_n_scalar; - uint32_t sv_n; - if (n_idx + ctx.s_block_stack > n_loop - 1) { - sv_n = ctx.cur_kv_seqlen - n_idx * pp_n_scalar; - } else { - sv_n = pp_n_scalar * ctx.s_block_stack; - } - params.sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; - - params.gm_split_idx = split_idx; - params.pp_n_scalar = pp_n_scalar; -} - -// 业务函数：TP1 Q 数据从 L1 加载到 L0A -// 复用非 TP1 的 PlatformWaitQLoadComplete / PlatformLoadQToL0A / PlatformSetQLoadComplete -// TP1 中 q_load_coeff = m，与非 TP1 一致 -__aicore__ __attribute__((always_inline)) inline void LoadTP1QDataToL0A( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformWaitQLoadComplete(embed_split_idx); - PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, m, false); - PlatformSetQLoadComplete(embed_split_idx); -} - -// 业务函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2） -// 更新 now_l1_offset 供后续 L1→L0B 使用 -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVMainFromGM( - TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 256; - PlatformLoadTP1KVMainToL1(embed_split_idx, params.qk_n, params.qk_round_n, - params.l1_kv_pingpong_flag, params.kv_offset, - embed_split_idx * 128); -} - -// 业务函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4） -// 更新 now_l1_offset 供后续 L1→L0B 使用 -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVRopeFromGM( - TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; - PlatformLoadTP1KVRopeToL1(params.qk_n, params.qk_round_n, - params.l1_kv_pingpong_flag, params.kv_offset_rope); -} - -// 业务函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVDataToL0B( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformLoadTP1KVToL0B(embed_split_idx, params.round_embed_split_size, - params.qk_round_n, params.now_l1_offset, - params.l1_kv_pingpong_flag); -} - -// 业务函数：TP1 QK MMAD 计算 -__aicore__ __attribute__((always_inline)) inline void ComputeTP1QKMMad( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformComputeTP1QKMMad(embed_split_idx, params.embed_split_size, - m, params.qk_n, params.l1_kv_pingpong_flag); -} - -// 业务函数：TP1 QK 结果拷贝到 GM（idx 4 时） -// GM 偏移: block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar -__aicore__ __attribute__((always_inline)) inline void CopyTP1QKResultToGM( - const TP1QKParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) -{ - uint64_t gm_dst_offset = (uint64_t)block_idx * TMP_SIZE_DECODER * 4 + - (uint64_t)((n_idx / 4) % 2) * TMP_SIZE_DECODER * 2 + - params.gm_split_idx * params.pp_n_scalar; - PlatformCopyTP1QKResultToGM(m, params.qk_round_n, - params.l1_kv_pingpong_flag, - gm_dst_offset, params.sv_round_n); -} // 业务函数：将 KV Rope 数据从 GM 搬运到 L1 // ND_FORMAT 场景：ND→NZ 格式转换（gm_to_l1），d_value=64, src_d_value=stride_kv_rope @@ -722,130 +474,6 @@ __aicore__ __attribute__((always_inline)) inline void ComputeQRope( } } -// ==================== TP1 QK 业务函数 ==================== - -// TP1 QK 参数结构体 -struct TP1QKParams { - uint32_t qk_n; - uint32_t qk_round_n; - uint32_t l1_kv_pingpong_flag; - uint32_t embed_split_size; - uint32_t round_embed_split_size; - int64_t now_l1_offset; - int64_t kv_offset; - int64_t kv_offset_rope; - uint32_t sv_round_n; - uint32_t gm_split_idx; // split_idx 用于 GM 偏移计算 - uint32_t pp_n_scalar; // block_size，用于 GM 偏移计算 -}; - -// 业务函数：初始化 TP1 QK 参数 -// 对应原始 CUBE1 stage1 中 split_idx 循环体内的参数计算 -__aicore__ __attribute__((always_inline)) inline void InitTP1QKParams( - TP1Context &ctx, uint32_t split_idx, uint32_t n_idx, TP1QKParams ¶ms) -{ - uint32_t pp_n_scalar = ctx.pp_n_scalar; - uint32_t n_loop = ctx.n_loop; - uint32_t now_idx = n_idx + split_idx; - - params.l1_kv_pingpong_flag = now_idx % 2; - - // 动态计算 qk_n / qk_round_n - if (now_idx == (n_loop - 1)) { - params.qk_n = (ctx.cur_kv_seqlen - now_idx * pp_n_scalar); - } else { - params.qk_n = pp_n_scalar; - } - params.qk_round_n = RoundUp(params.qk_n); - - // embed_split_size 默认 128，idx==4 时为 64（在编排层设置） - params.embed_split_size = 128; - params.round_embed_split_size = RoundUp(params.embed_split_size); - - // now_l1_offset 默认 0，在 LoadTP1KVMain/KVRope 中更新 - params.now_l1_offset = 0; - - // block_table_id / kv_offset / kv_offset_rope 地址计算 - uint32_t block_table_id = (uint32_t)(*(block_tables_gm + - ctx.cur_batch * max_num_blocks_per_query + ctx.start_kv / block_size + now_idx)); - params.kv_offset = (int64_t)block_table_id * block_size * stride_kv; - params.kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; - - // sv_round_n 计算（用于 S→GM 的 dstStride） - uint32_t sv_n_triu = n_loop * pp_n_scalar; - uint32_t sv_n; - if (n_idx + ctx.s_block_stack > n_loop - 1) { - sv_n = ctx.cur_kv_seqlen - n_idx * pp_n_scalar; - } else { - sv_n = pp_n_scalar * ctx.s_block_stack; - } - params.sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; - - params.gm_split_idx = split_idx; - params.pp_n_scalar = pp_n_scalar; -} - -// 业务函数：TP1 Q 数据从 L1 加载到 L0A -// 复用非 TP1 的 PlatformWaitQLoadComplete / PlatformLoadQToL0A / PlatformSetQLoadComplete -// TP1 中 q_load_coeff = m，与非 TP1 一致 -__aicore__ __attribute__((always_inline)) inline void LoadTP1QDataToL0A( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformWaitQLoadComplete(embed_split_idx); - PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, m, false); - PlatformSetQLoadComplete(embed_split_idx); -} - -// 业务函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2） -// 更新 now_l1_offset 供后续 L1→L0B 使用 -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVMainFromGM( - TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 256; - PlatformLoadTP1KVMainToL1(embed_split_idx, params.qk_n, params.qk_round_n, - params.l1_kv_pingpong_flag, params.kv_offset, - embed_split_idx * 128); -} - -// 业务函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4） -// 更新 now_l1_offset 供后续 L1→L0B 使用 -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVRopeFromGM( - TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; - PlatformLoadTP1KVRopeToL1(params.qk_n, params.qk_round_n, - params.l1_kv_pingpong_flag, params.kv_offset_rope); -} - -// 业务函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVDataToL0B( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformLoadTP1KVToL0B(embed_split_idx, params.round_embed_split_size, - params.qk_round_n, params.now_l1_offset, - params.l1_kv_pingpong_flag); -} - -// 业务函数：TP1 QK MMAD 计算 -__aicore__ __attribute__((always_inline)) inline void ComputeTP1QKMMad( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformComputeTP1QKMMad(embed_split_idx, params.embed_split_size, - m, params.qk_n, params.l1_kv_pingpong_flag); -} - -// 业务函数：TP1 QK 结果拷贝到 GM（idx 4 时） -// GM 偏移: block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar -__aicore__ __attribute__((always_inline)) inline void CopyTP1QKResultToGM( - const TP1QKParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) -{ - uint64_t gm_dst_offset = (uint64_t)block_idx * TMP_SIZE_DECODER * 4 + - (uint64_t)((n_idx / 4) % 2) * TMP_SIZE_DECODER * 2 + - params.gm_split_idx * params.pp_n_scalar; - PlatformCopyTP1QKResultToGM(m, params.qk_round_n, - params.l1_kv_pingpong_flag, - gm_dst_offset, params.sv_round_n); -} // 业务函数：调度 Cube MLA 任务（Run 方法的 for 循环部分） // 遍历所有 batch，解析 tiling 参数，调用 InnerRunCubeMLA @@ -879,130 +507,6 @@ __aicore__ __attribute__((always_inline)) inline void ScheduleCubeTasks() } } -// ==================== TP1 QK 业务函数 ==================== - -// TP1 QK 参数结构体 -struct TP1QKParams { - uint32_t qk_n; - uint32_t qk_round_n; - uint32_t l1_kv_pingpong_flag; - uint32_t embed_split_size; - uint32_t round_embed_split_size; - int64_t now_l1_offset; - int64_t kv_offset; - int64_t kv_offset_rope; - uint32_t sv_round_n; - uint32_t gm_split_idx; // split_idx 用于 GM 偏移计算 - uint32_t pp_n_scalar; // block_size，用于 GM 偏移计算 -}; - -// 业务函数：初始化 TP1 QK 参数 -// 对应原始 CUBE1 stage1 中 split_idx 循环体内的参数计算 -__aicore__ __attribute__((always_inline)) inline void InitTP1QKParams( - TP1Context &ctx, uint32_t split_idx, uint32_t n_idx, TP1QKParams ¶ms) -{ - uint32_t pp_n_scalar = ctx.pp_n_scalar; - uint32_t n_loop = ctx.n_loop; - uint32_t now_idx = n_idx + split_idx; - - params.l1_kv_pingpong_flag = now_idx % 2; - - // 动态计算 qk_n / qk_round_n - if (now_idx == (n_loop - 1)) { - params.qk_n = (ctx.cur_kv_seqlen - now_idx * pp_n_scalar); - } else { - params.qk_n = pp_n_scalar; - } - params.qk_round_n = RoundUp(params.qk_n); - - // embed_split_size 默认 128，idx==4 时为 64（在编排层设置） - params.embed_split_size = 128; - params.round_embed_split_size = RoundUp(params.embed_split_size); - - // now_l1_offset 默认 0，在 LoadTP1KVMain/KVRope 中更新 - params.now_l1_offset = 0; - - // block_table_id / kv_offset / kv_offset_rope 地址计算 - uint32_t block_table_id = (uint32_t)(*(block_tables_gm + - ctx.cur_batch * max_num_blocks_per_query + ctx.start_kv / block_size + now_idx)); - params.kv_offset = (int64_t)block_table_id * block_size * stride_kv; - params.kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; - - // sv_round_n 计算（用于 S→GM 的 dstStride） - uint32_t sv_n_triu = n_loop * pp_n_scalar; - uint32_t sv_n; - if (n_idx + ctx.s_block_stack > n_loop - 1) { - sv_n = ctx.cur_kv_seqlen - n_idx * pp_n_scalar; - } else { - sv_n = pp_n_scalar * ctx.s_block_stack; - } - params.sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; - - params.gm_split_idx = split_idx; - params.pp_n_scalar = pp_n_scalar; -} - -// 业务函数：TP1 Q 数据从 L1 加载到 L0A -// 复用非 TP1 的 PlatformWaitQLoadComplete / PlatformLoadQToL0A / PlatformSetQLoadComplete -// TP1 中 q_load_coeff = m，与非 TP1 一致 -__aicore__ __attribute__((always_inline)) inline void LoadTP1QDataToL0A( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformWaitQLoadComplete(embed_split_idx); - PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, m, false); - PlatformSetQLoadComplete(embed_split_idx); -} - -// 业务函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2） -// 更新 now_l1_offset 供后续 L1→L0B 使用 -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVMainFromGM( - TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 256; - PlatformLoadTP1KVMainToL1(embed_split_idx, params.qk_n, params.qk_round_n, - params.l1_kv_pingpong_flag, params.kv_offset, - embed_split_idx * 128); -} - -// 业务函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4） -// 更新 now_l1_offset 供后续 L1→L0B 使用 -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVRopeFromGM( - TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; - PlatformLoadTP1KVRopeToL1(params.qk_n, params.qk_round_n, - params.l1_kv_pingpong_flag, params.kv_offset_rope); -} - -// 业务函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVDataToL0B( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformLoadTP1KVToL0B(embed_split_idx, params.round_embed_split_size, - params.qk_round_n, params.now_l1_offset, - params.l1_kv_pingpong_flag); -} - -// 业务函数：TP1 QK MMAD 计算 -__aicore__ __attribute__((always_inline)) inline void ComputeTP1QKMMad( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformComputeTP1QKMMad(embed_split_idx, params.embed_split_size, - m, params.qk_n, params.l1_kv_pingpong_flag); -} - -// 业务函数：TP1 QK 结果拷贝到 GM（idx 4 时） -// GM 偏移: block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar -__aicore__ __attribute__((always_inline)) inline void CopyTP1QKResultToGM( - const TP1QKParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) -{ - uint64_t gm_dst_offset = (uint64_t)block_idx * TMP_SIZE_DECODER * 4 + - (uint64_t)((n_idx / 4) % 2) * TMP_SIZE_DECODER * 2 + - params.gm_split_idx * params.pp_n_scalar; - PlatformCopyTP1QKResultToGM(m, params.qk_round_n, - params.l1_kv_pingpong_flag, - gm_dst_offset, params.sv_round_n); -} // 业务函数：调度 Cube MLA TP1 任务（RunTP1 方法的中段业务逻辑） // 包含主循环调度 + tail 优化三分支（cores_per_seq 动态调整） @@ -1190,128 +694,3 @@ __aicore__ __attribute__((always_inline)) inline void ScheduleCubeTasksTP1() } } } - -// ==================== TP1 QK 业务函数 ==================== - -// TP1 QK 参数结构体 -struct TP1QKParams { - uint32_t qk_n; - uint32_t qk_round_n; - uint32_t l1_kv_pingpong_flag; - uint32_t embed_split_size; - uint32_t round_embed_split_size; - int64_t now_l1_offset; - int64_t kv_offset; - int64_t kv_offset_rope; - uint32_t sv_round_n; - uint32_t gm_split_idx; // split_idx 用于 GM 偏移计算 - uint32_t pp_n_scalar; // block_size，用于 GM 偏移计算 -}; - -// 业务函数：初始化 TP1 QK 参数 -// 对应原始 CUBE1 stage1 中 split_idx 循环体内的参数计算 -__aicore__ __attribute__((always_inline)) inline void InitTP1QKParams( - TP1Context &ctx, uint32_t split_idx, uint32_t n_idx, TP1QKParams ¶ms) -{ - uint32_t pp_n_scalar = ctx.pp_n_scalar; - uint32_t n_loop = ctx.n_loop; - uint32_t now_idx = n_idx + split_idx; - - params.l1_kv_pingpong_flag = now_idx % 2; - - // 动态计算 qk_n / qk_round_n - if (now_idx == (n_loop - 1)) { - params.qk_n = (ctx.cur_kv_seqlen - now_idx * pp_n_scalar); - } else { - params.qk_n = pp_n_scalar; - } - params.qk_round_n = RoundUp(params.qk_n); - - // embed_split_size 默认 128，idx==4 时为 64（在编排层设置） - params.embed_split_size = 128; - params.round_embed_split_size = RoundUp(params.embed_split_size); - - // now_l1_offset 默认 0，在 LoadTP1KVMain/KVRope 中更新 - params.now_l1_offset = 0; - - // block_table_id / kv_offset / kv_offset_rope 地址计算 - uint32_t block_table_id = (uint32_t)(*(block_tables_gm + - ctx.cur_batch * max_num_blocks_per_query + ctx.start_kv / block_size + now_idx)); - params.kv_offset = (int64_t)block_table_id * block_size * stride_kv; - params.kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; - - // sv_round_n 计算（用于 S→GM 的 dstStride） - uint32_t sv_n_triu = n_loop * pp_n_scalar; - uint32_t sv_n; - if (n_idx + ctx.s_block_stack > n_loop - 1) { - sv_n = ctx.cur_kv_seqlen - n_idx * pp_n_scalar; - } else { - sv_n = pp_n_scalar * ctx.s_block_stack; - } - params.sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; - - params.gm_split_idx = split_idx; - params.pp_n_scalar = pp_n_scalar; -} - -// 业务函数：TP1 Q 数据从 L1 加载到 L0A -// 复用非 TP1 的 PlatformWaitQLoadComplete / PlatformLoadQToL0A / PlatformSetQLoadComplete -// TP1 中 q_load_coeff = m，与非 TP1 一致 -__aicore__ __attribute__((always_inline)) inline void LoadTP1QDataToL0A( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformWaitQLoadComplete(embed_split_idx); - PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, m, false); - PlatformSetQLoadComplete(embed_split_idx); -} - -// 业务函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2） -// 更新 now_l1_offset 供后续 L1→L0B 使用 -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVMainFromGM( - TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 256; - PlatformLoadTP1KVMainToL1(embed_split_idx, params.qk_n, params.qk_round_n, - params.l1_kv_pingpong_flag, params.kv_offset, - embed_split_idx * 128); -} - -// 业务函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4） -// 更新 now_l1_offset 供后续 L1→L0B 使用 -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVRopeFromGM( - TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; - PlatformLoadTP1KVRopeToL1(params.qk_n, params.qk_round_n, - params.l1_kv_pingpong_flag, params.kv_offset_rope); -} - -// 业务函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVDataToL0B( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformLoadTP1KVToL0B(embed_split_idx, params.round_embed_split_size, - params.qk_round_n, params.now_l1_offset, - params.l1_kv_pingpong_flag); -} - -// 业务函数：TP1 QK MMAD 计算 -__aicore__ __attribute__((always_inline)) inline void ComputeTP1QKMMad( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformComputeTP1QKMMad(embed_split_idx, params.embed_split_size, - m, params.qk_n, params.l1_kv_pingpong_flag); -} - -// 业务函数：TP1 QK 结果拷贝到 GM（idx 4 时） -// GM 偏移: block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar -__aicore__ __attribute__((always_inline)) inline void CopyTP1QKResultToGM( - const TP1QKParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) -{ - uint64_t gm_dst_offset = (uint64_t)block_idx * TMP_SIZE_DECODER * 4 + - (uint64_t)((n_idx / 4) % 2) * TMP_SIZE_DECODER * 2 + - params.gm_split_idx * params.pp_n_scalar; - PlatformCopyTP1QKResultToGM(m, params.qk_round_n, - params.l1_kv_pingpong_flag, - gm_dst_offset, params.sv_round_n); -} \ No newline at end of file From bd2912af1237e604d03c7a7444898d97e215ae60 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Fri, 14 Aug 2026 16:25:32 +0800 Subject: [PATCH 19/24] feat: adapter the mla a5 --- .../include/iterators/l0c_to_gm_iterator.inc | 11 +- .../include/iterators/l1_to_l0_iterator.inc | 19 + .../op_kernel/mixkernels/include/mem.h | 2 +- .../op_kernel/mixkernels/include/mma.h | 19 +- .../op_kernel/mixkernels/include/simd.h | 11 + .../op_kernel/multi_latent_attention.cpp | 51 +- .../op_kernel/multi_latent_attention.h | 3716 ++++++++++++++--- .../multi_latent_attention_aiv_arch32.h | 1089 ----- .../op_kernel/multi_latent_attention_aiv_bs.h | 1303 ------ .../op_kernel/multi_latent_attention_arch32.h | 640 --- .../op_kernel/multi_latent_attention_bs.h | 696 --- 11 files changed, 3296 insertions(+), 4261 deletions(-) delete mode 100644 xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h delete mode 100644 xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h delete mode 100644 xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_arch32.h delete mode 100644 xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h diff --git a/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/iterators/l0c_to_gm_iterator.inc b/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/iterators/l0c_to_gm_iterator.inc index 6a7f4f7..b14536b 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/iterators/l0c_to_gm_iterator.inc +++ b/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/iterators/l0c_to_gm_iterator.inc @@ -23,13 +23,18 @@ struct l0c_to_gm { uint32_t dstStride) { #if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) - // A5 architecture: use FixpipeParamsArch3510, F322F16 mode does not require cbufWorkspace + // A5 architecture: use FixpipeParamsArch3510, F322F16 mode. + // On 3510, L0C srcStride = RoundUp<16>(m) which matches the official + // arch35 reference (FixpipeOut.h). AscendC::FixpipeParamsArch3510 intriParams( nTileActual, // nSize mTileActual, // mSize srcStride, // srcStride dstStride); // dstStride intriParams.quantPre = QuantMode_t::F322F16; + intriParams.params.ndNum = 1; + intriParams.params.srcNdStride = 0; + intriParams.params.dstNdStride = 0; AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); #elif defined(__DAV_C220_CUBE__) // V220 architecture (A2/A3): use FixpipeParamsV220 @@ -112,7 +117,7 @@ struct l0c_to_gm { uint32_t dstStride) { #if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) - // A5 architecture: use FixpipeParamsArch3510, F322BF16 mode + // A5 architecture: use FixpipeParamsArch3510, F322BF16 mode. AscendC::FixpipeParamsArch3510 intriParams( nTileActual, // nSize mTileActual, // mSize @@ -155,7 +160,7 @@ struct l0c_to_gm { uint32_t dstStride) { #if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) - // A5 architecture: use FixpipeParamsArch3510, NoQuant mode + // A5 architecture: use FixpipeParamsArch3510, NoQuant mode. AscendC::FixpipeParamsArch3510 intriParams( nTileActual, // nSize mTileActual, // mSize diff --git a/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/iterators/l1_to_l0_iterator.inc b/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/iterators/l1_to_l0_iterator.inc index 4a961e2..ec2be6a 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/iterators/l1_to_l0_iterator.inc +++ b/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/iterators/l1_to_l0_iterator.inc @@ -9,6 +9,7 @@ template struct l1_to_l0_a { using HardwareParams = HardwareInfo; static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + static constexpr uint32_t K0 = 16; // elements per K-direction fractal (C0_SIZE for half) __aicore__ l1_to_l0_a(AscendC::LocalTensor l0Tensor, AscendC::LocalTensor l1Tensor, @@ -19,6 +20,23 @@ struct l1_to_l0_a(kSrcStride), // mStep + static_cast(kPartCeil), // kStep + static_cast(kSrcStride), // srcStride + static_cast(kSrcStride), // dstStride + IsTransPose, // ifTranspose + 0)); // sid +#else AscendC::LoadData(l0Tensor, l1Tensor, AscendC::LoadData2dParams(0, // baseIdx @@ -28,6 +46,7 @@ struct l1_to_l0_a(AscendC::TPosition::VECIN); -#elif __DAV_C220_CUBE__ +#elif defined(__DAV_C220_CUBE__) tensor[(uint32_t)BufferType::ASCEND_CB].InitBuffer(0, bufferSize[(uint32_t)BufferType::ASCEND_CB]); tensor[(uint32_t)BufferType::ASCEND_CB].address_.logicPos = static_cast(AscendC::TPosition::A1); tensor[(uint32_t)BufferType::ASCEND_L0A].InitBuffer(0, bufferSize[(uint32_t)BufferType::ASCEND_L0A]); diff --git a/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/mma.h b/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/mma.h index 1e6550a..091e680 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/mma.h +++ b/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/mma.h @@ -35,10 +35,11 @@ struct mmad { uint32_t kPartActual, bool initC) { - AscendC::Mmad(l0cTensor, - l0aTensor, - l0bTensor, - AscendC::MmadParams(mTileActual, nTileActual, kPartActual, 0, false, initC)); + AscendC::MmadParams mmadParams(mTileActual, nTileActual, kPartActual, 0, false, initC); +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) + mmadParams.disableGemv = true; +#endif + AscendC::Mmad(l0cTensor, l0aTensor, l0bTensor, mmadParams); }; __aicore__ mmad(AscendC::LocalTensor l0cTensor, @@ -53,11 +54,11 @@ struct mmad { 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)); + AscendC::MmadParams mmadParams(mTileActual, nTileActual, kPartActual, 0, false, initC); +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) + mmadParams.disableGemv = true; +#endif + AscendC::Mmad(l0cTensor, l0aTensor, l0bTensor, biasTensor, mmadParams); }; }; diff --git a/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/simd.h b/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/simd.h index 76c0b57..3947bd9 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/simd.h +++ b/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/simd.h @@ -186,12 +186,23 @@ __aicore__ inline void exp_v(AscendC::LocalTensor dst, uint16_t dstRepeatStride, uint16_t srcRepeatStride) { +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) + // 3510: Use PRECISION_1ULP_FTZ_FALSE to preserve Subnormal numbers, + // improving online softmax precision for large kv_seqlen. + AscendC::Exp( + dst, + src, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +#else AscendC::Exp( dst, src, (uint64_t)0, repeat, AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +#endif } ///////////////////////////////////////////////////// diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.cpp b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.cpp index bfc9f23..5d87ab2 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.cpp +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.cpp @@ -13,6 +13,19 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +// 3510 (Ascend950/A5) arch mapping: __DAV_C220_CUBE__/__DAV_C220_VEC__ are +// compiler-predefined on 2201 (A3) but NOT on 3510 (A5). On 3510 the compiler +// defines __DAV_CUBE__ (cube core) / __DAV_VEC__ (vector core) instead. +// Map them here so the existing V220 code paths compile and execute on 3510. +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) + #if defined(__DAV_CUBE__) && !defined(__DAV_C220_CUBE__) + #define __DAV_C220_CUBE__ + #endif + #if defined(__DAV_VEC__) && !defined(__DAV_C220_VEC__) + #define __DAV_C220_VEC__ + #endif +#endif + #include "kernel_operator.h" #include "multi_latent_attention.h" #include "lib/matmul_intf.h" @@ -48,7 +61,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A SetMasknorm(); #ifdef __DAV_C220_VEC__ SetVectorMask((uint64_t)-1, (uint64_t)-1); -#elif __DAV_C220_CUBE__ +#elif defined(__DAV_C220_CUBE__) SetPadding(0); SetNdpara(1, 0, 0); #endif @@ -57,7 +70,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_fp16 {}; pa_aic_fp16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_fp16.Run(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.Run(); @@ -67,7 +80,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_bf16 {}; pa_aic_bf16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_bf16.Run(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.Run(); @@ -77,7 +90,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_fp16 {}; pa_aic_fp16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_fp16.Run(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.Run(); @@ -87,7 +100,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_bf16 {}; pa_aic_bf16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_bf16.Run(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.Run(); @@ -97,7 +110,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_fp16 {}; pa_aic_fp16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_fp16.RunTP1(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.RunTP1(); @@ -107,7 +120,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_bf16 {}; pa_aic_bf16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_bf16.RunTP1(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.RunTP1(); @@ -117,7 +130,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_fp16 {}; pa_aic_fp16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_fp16.RunTP1(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.RunTP1(); @@ -127,7 +140,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_bf16 {}; pa_aic_bf16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_bf16.RunTP1(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.RunTP1(); @@ -137,7 +150,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_fp16 {}; pa_aic_fp16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_fp16.Run(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.SetArgs2(lse_gm); @@ -148,7 +161,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_bf16 {}; pa_aic_bf16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_bf16.Run(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.SetArgs2(lse_gm); @@ -159,7 +172,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_fp16 {}; pa_aic_fp16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_fp16.Run(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.SetArgs2(lse_gm); @@ -170,7 +183,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_bf16 {}; pa_aic_bf16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_bf16.Run(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.SetArgs2(lse_gm); @@ -181,7 +194,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_fp16 {}; pa_aic_fp16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_fp16.RunTP1(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.SetArgs2(lse_gm); @@ -192,7 +205,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_fp16 {}; pa_aic_fp16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_fp16.RunTP1(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.SetArgs2(lse_gm); @@ -203,7 +216,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_fp16 {}; pa_aic_fp16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_fp16.RunTP1(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.SetArgs2(lse_gm); @@ -214,7 +227,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_fp16 {}; pa_aic_fp16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_fp16.RunTP1(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.SetArgs2(lse_gm); @@ -225,7 +238,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_fp16 {}; pa_aic_fp16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_fp16.Run(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.Run(); @@ -235,7 +248,7 @@ extern "C" __global__ __aicore__ void multi_latent_attention(GM_ADDR query, GM_A MLAttentionDecoderAic pa_aic_bf16 {}; pa_aic_bf16.SetArgs(q_gm, q_rope_gm, ctkv_gm, ctkv_rope_gm, block_tables_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, tiling_para_gm); pa_aic_bf16.Run(); -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) MLADecoderAiv pa_aiv {}; pa_aiv.SetArgs(block_tables_gm, deq_qk_gm, deq_pv_gm, o_gm, s_gm, s_rope_out_gm, p_gm, o_tmp_gm, go_gm, tmp_gm, tiling_para_gm, mask_gm); pa_aiv.Run(); diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h index af4d63b..6e46db9 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h +++ b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention.h @@ -163,7 +163,7 @@ constexpr uint64_t CONST_128 = 128; constexpr uint32_t EMBED_SPLIT = 256; constexpr uint32_t ROUND_EMBED_SPLIT = 256; -#elif __DAV_C220_VEC__ +#elif defined(__DAV_C220_VEC__) constexpr uint32_t HALF_VECTOR_SIZE = 128; constexpr uint32_t UB_ALIGN_BYTE = 32; constexpr int64_t UB_UINT8_BLOCK_SIZE_MLA = 16384; // 96 * 128 * 2B // prefill/decoder diff @@ -368,443 +368,1167 @@ class MLAttentionDecoderAic { __aicore__ __attribute__((always_inline)) inline void Run() { - PlatformInitPipeSync(); - ScheduleCubeTasks(); - PlatformWaitPipeSync(); - } - - __aicore__ __attribute__((always_inline)) inline void RunTP1() - { - PlatformInitPipeSync(); - ScheduleCubeTasksTP1(); - PlatformWaitPipeSync(); - } - -private: - // ====== MLA Cube refactor: context struct + sub-functions ====== - - struct MLAContext { - // input params - uint32_t cur_batch; - uint32_t start_head; - uint32_t cur_head_num; - uint32_t start_kv; - uint32_t cur_q_seqlen; - uint32_t cur_kv_seqlen; - uint32_t offset_tiling; - - // Q address - uint64_t q_offset; - uint64_t q_rope_offset; - - // loop & size - uint32_t pp_n_scalar; - uint32_t sub_n_loop; - uint32_t n_loop; - - // QK dims - uint32_t qk_n; - uint32_t qk_round_n; - uint32_t qk_n_2; - uint32_t qk_round_n_2; - uint32_t qk_round_n_l1; - uint32_t qk_round_n_2_l1; - - // hidden size - uint64_t hidden_size; - - // K round - uint64_t k_round_n; - - // row info - uint32_t row_num; - // m is class member - }; - - __aicore__ __attribute__((always_inline)) inline void InitMLAContext( - MLAContext &ctx, uint32_t cur_batch, uint32_t start_head, uint32_t cur_head_num, - uint32_t start_kv, uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t offset_tiling) - { - ctx.cur_batch = cur_batch; - ctx.start_head = start_head; - ctx.cur_head_num = cur_head_num; - ctx.start_kv = start_kv; - ctx.cur_q_seqlen = cur_q_seqlen; - ctx.cur_kv_seqlen = cur_kv_seqlen; - ctx.offset_tiling = offset_tiling; - - uint32_t addr_q_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); - uint32_t addr_q_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 3 + offset_tiling)); - uint64_t addr_q_scalar = (uint64_t)(((uint64_t)addr_q_high32) << 32 | addr_q_loww32); - ctx.q_offset = addr_q_scalar * 512 + start_head * 512; - ctx.q_rope_offset = addr_q_scalar * 64 + start_head * 64; - - ctx.pp_n_scalar = block_size; - ctx.sub_n_loop = ctx.pp_n_scalar / block_size; - ctx.n_loop = (cur_kv_seqlen + ctx.pp_n_scalar - 1) / ctx.pp_n_scalar; + SET_FLAG(M, MTE1, EVENT_ID0); + SET_FLAG(M, MTE1, EVENT_ID1); + SET_FLAG(M, MTE1, EVENT_ID2); + SET_FLAG(M, MTE1, EVENT_ID3); + SET_FLAG(M, MTE1, EVENT_ID4); + SET_FLAG(M, MTE1, EVENT_ID5); + SET_FLAG(M, MTE1, EVENT_ID6); + SET_FLAG(M, MTE1, EVENT_ID7); + SET_FLAG(FIX, M, EVENT_ID0); + SET_FLAG(FIX, M, EVENT_ID1); + SET_FLAG(MTE1, MTE2, EVENT_ID0); + SET_FLAG(MTE1, MTE2, EVENT_ID1); + SET_FLAG(MTE1, MTE2, EVENT_ID2); + SET_FLAG(MTE1, MTE2, EVENT_ID3); + SET_FLAG(MTE1, MTE2, EVENT_ID4); + SET_FLAG(MTE1, MTE2, EVENT_ID5); + SET_FLAG(MTE1, MTE2, EVENT_ID6); + SET_FLAG(MTE1, MTE2, EVENT_ID7); + SET_FLAG(FIX, MTE1, EVENT_ID0); + SET_FLAG(FIX, MTE1, EVENT_ID1); + SET_FLAG(FIX, MTE1, EVENT_ID2); + SET_FLAG(FIX, MTE1, EVENT_ID3); + SET_FLAG(FIX, MTE1, EVENT_ID4); + SET_FLAG(FIX, MTE1, EVENT_ID5); + SET_FLAG(MTE2, FIX, EVENT_ID0); + + + uint64_t cur_batch = 0; + + uint32_t q_block_num_per_batch = (q_heads + cur_qn_blk_size - 1) / cur_qn_blk_size; + uint32_t process_num = q_block_num_per_batch * num_batches; + + for (uint32_t process = block_idx; process < process_num; process += (uint32_t)block_num) { // for task + cur_batch = process / q_block_num_per_batch; + if (cur_batch >= num_batches) break; + + uint32_t offset_tiling = tiling_head_size + tiling_para_size * cur_batch; + uint32_t start_core_idx = (cur_batch * q_block_num_per_batch) % block_num; + + uint32_t q_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); + if (kv_seqlen == 0) { + continue; + } + uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; - ctx.qk_n = ctx.pp_n_scalar; - ctx.qk_round_n = RoundUp(ctx.qk_n); - ctx.qk_n_2 = ctx.pp_n_scalar; - ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); - ctx.qk_round_n_l1 = RoundUp(ctx.qk_n); - ctx.qk_round_n_2_l1 = RoundUp(ctx.qk_n_2); + uint32_t start_head = (process % q_block_num_per_batch) * cur_qn_blk_size; + uint32_t start_kv = 0; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen; + uint32_t cur_head_num = cur_qn_blk_size; - ctx.hidden_size = 576; - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - ctx.hidden_size = 512; + InnerRunCubeMLA(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, + offset_tiling); } - ctx.k_round_n = ctx.qk_round_n; - - ctx.row_num = cur_head_num * cur_q_seqlen; - m = RoundUp<16>(ctx.row_num); - } - - // === 第一层：LoadQData 编排层 === - // 业务函数（LoadQMainFromGMToL1 / LoadQRopeFromGMToL1 / PlatformSetQLoadComplete） - // 已迁移至 bs.h / arch32.h，通过 include 展开 - - // 业务函数：加载 Q 数据（主体 + Rope）从 GM 到 L1 - __aicore__ __attribute__((always_inline)) inline void LoadQData(MLAContext &ctx) - { - uint32_t cur_q_seqlen = ctx.cur_q_seqlen; - uint32_t cur_head_num = ctx.cur_head_num; - uint64_t q_offset = ctx.q_offset; - uint64_t q_rope_offset = ctx.q_rope_offset; - - // copy Q - LoadQMainFromGMToL1(l1q_buf_addr_tensor, q_gm_tensor, q_offset, cur_q_seqlen, cur_head_num); - LoadQRopeFromGMToL1(l1q_buf_addr_tensor, l1q_rope_buf_addr_tensor, q_rope_gm_tensor, q_rope_offset, cur_q_seqlen, cur_head_num); - PlatformSetQLoadComplete(); + WAIT_FLAG(M, MTE1, EVENT_ID0); + WAIT_FLAG(M, MTE1, EVENT_ID1); + WAIT_FLAG(M, MTE1, EVENT_ID2); + WAIT_FLAG(M, MTE1, EVENT_ID3); + WAIT_FLAG(M, MTE1, EVENT_ID4); + WAIT_FLAG(M, MTE1, EVENT_ID5); + WAIT_FLAG(M, MTE1, EVENT_ID6); + WAIT_FLAG(M, MTE1, EVENT_ID7); + WAIT_FLAG(FIX, M, EVENT_ID0); + WAIT_FLAG(FIX, M, EVENT_ID1); + WAIT_FLAG(MTE1, MTE2, EVENT_ID0); + WAIT_FLAG(MTE1, MTE2, EVENT_ID1); + WAIT_FLAG(MTE1, MTE2, EVENT_ID2); + WAIT_FLAG(MTE1, MTE2, EVENT_ID3); + WAIT_FLAG(MTE1, MTE2, EVENT_ID4); + WAIT_FLAG(MTE1, MTE2, EVENT_ID5); + WAIT_FLAG(MTE1, MTE2, EVENT_ID6); + WAIT_FLAG(MTE1, MTE2, EVENT_ID7); + WAIT_FLAG(FIX, MTE1, EVENT_ID0); + WAIT_FLAG(FIX, MTE1, EVENT_ID1); + WAIT_FLAG(FIX, MTE1, EVENT_ID2); + WAIT_FLAG(FIX, MTE1, EVENT_ID3); + WAIT_FLAG(FIX, MTE1, EVENT_ID4); + WAIT_FLAG(FIX, MTE1, EVENT_ID5); + WAIT_FLAG(MTE2, FIX, EVENT_ID0); + PIPE_BARRIER(ALL); } - // === 第一层：LoadKVData 编排层 === - // 业务函数已迁移至 bs.h / arch32.h - - __aicore__ __attribute__((always_inline)) inline void LoadKVData( - MLAContext &ctx, uint32_t n_idx, uint32_t l1_kv_pingpong_flag) + __aicore__ __attribute__((always_inline)) inline void RunTP1() { - uint32_t qk_n = ctx.qk_n; - uint32_t qk_round_n = ctx.qk_round_n; - uint32_t qk_round_n_l1 = ctx.qk_round_n_l1; - uint32_t cur_batch = ctx.cur_batch; - uint32_t start_kv = ctx.start_kv; - - uint32_t block_table_id = (uint32_t)(*(block_tables_gm + - cur_batch * max_num_blocks_per_query + start_kv / block_size + n_idx)); - int64_t kv_offset = (int64_t)block_table_id * block_size * stride_kv; - int64_t kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; - - PlatformWaitKVLoadReady(l1_kv_pingpong_flag); - if constexpr (KInputType == InputFormat::ND_FORMAT) { - // 分支1: ND→NZ，K 主体搬到 l1kv_buf，Rope 紧随其后 - LoadKVMainFromGMToL1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576], - k_gm_tensor[kv_offset], - qk_n, qk_round_n, 512, stride_kv, true /*ND→NZ*/); - PlatformSetKVMainLoadComplete(l1_kv_pingpong_flag); - LoadKVRopeFromGMToL1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576 + 512 * qk_round_n], - k_rope_gm_tensor[kv_offset_rope], - qk_n, qk_round_n, 64, stride_kv_rope, true /*ND→NZ*/); - } else if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - // 分支2: NZ→NZ，K 主体搬到 l1kv_buf(512布局)，Rope 搬到独立 l1kv_rope_buf - LoadKVMainFromGMToL1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 512], - k_gm_tensor[kv_offset], - qk_round_n_l1, block_size, 512, 0, false /*NZ→NZ*/); - PlatformSetKVMainLoadComplete(l1_kv_pingpong_flag); - LoadKVRopeFromGMToL1( - l1kv_rope_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64], - k_rope_gm_tensor[kv_offset_rope], - qk_round_n, block_size, 64, 0, false /*NZ→NZ*/); - } else { - // 分支3: NZ→NZ，K 主体搬到 l1kv_buf(576布局)，Rope 紧随其后 - LoadKVMainFromGMToL1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576], - k_gm_tensor[kv_offset], - qk_round_n, block_size, 512, 0, false /*NZ→NZ*/); - PlatformSetKVMainLoadComplete(l1_kv_pingpong_flag); - LoadKVRopeFromGMToL1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576 + 512 * qk_round_n], - k_rope_gm_tensor[kv_offset_rope], - qk_round_n, block_size, 64, 0, false /*NZ→NZ*/); - } + SET_FLAG(M, MTE1, EVENT_ID0); + SET_FLAG(M, MTE1, EVENT_ID1); + SET_FLAG(M, MTE1, EVENT_ID2); + SET_FLAG(M, MTE1, EVENT_ID3); + SET_FLAG(M, MTE1, EVENT_ID4); + SET_FLAG(M, MTE1, EVENT_ID5); + SET_FLAG(M, MTE1, EVENT_ID6); + SET_FLAG(M, MTE1, EVENT_ID7); + SET_FLAG(FIX, M, EVENT_ID0); + SET_FLAG(FIX, M, EVENT_ID1); + SET_FLAG(MTE1, MTE2, EVENT_ID0); + SET_FLAG(MTE1, MTE2, EVENT_ID1); + SET_FLAG(MTE1, MTE2, EVENT_ID2); + SET_FLAG(MTE1, MTE2, EVENT_ID3); + SET_FLAG(MTE1, MTE2, EVENT_ID4); + SET_FLAG(MTE1, MTE2, EVENT_ID5); + SET_FLAG(MTE1, MTE2, EVENT_ID6); + SET_FLAG(MTE1, MTE2, EVENT_ID7); + SET_FLAG(FIX, MTE1, EVENT_ID0); + SET_FLAG(FIX, MTE1, EVENT_ID1); + SET_FLAG(FIX, MTE1, EVENT_ID2); + SET_FLAG(FIX, MTE1, EVENT_ID3); + SET_FLAG(FIX, MTE1, EVENT_ID4); + SET_FLAG(FIX, MTE1, EVENT_ID5); + SET_FLAG(MTE2, FIX, EVENT_ID0); - PlatformSetKVRopeLoadComplete(l1_kv_pingpong_flag); - } + uint32_t tail = totalTaskNum % block_num; + if constexpr (EnableOptimization) { - // ==================== ComputeQK 三层架构拆分 ==================== - // === 第一层：QKParams 结构体 === - struct QKParams { - uint32_t qk_n; - uint32_t qk_round_n; - uint32_t qk_round_n_l1; - uint64_t hidden_size; - uint64_t k_round_n; - uint32_t row_num; - uint32_t embed_split_size; - uint32_t round_embed_split_size; - uint32_t q_load_coeff; - uint64_t hidden_split_time; - uint32_t l1_kv_pingpong_flag; - }; - - // [QK 第三层平台函数 + 第二层业务函数] 已迁移至 arch32.h / bs.h - // （include 指令统一放在 PVParams 结构体之后，确保所有 Params 类型可见） - - // === 第一层：ComputeQK 编排层 === - __aicore__ __attribute__((always_inline)) inline void ComputeQK( - MLAContext &ctx, uint32_t n_idx, uint32_t l1_kv_pingpong_flag) - { - QKParams params; - InitQKParams(ctx, l1_kv_pingpong_flag, params); - - for (uint64_t embed_split_idx = 0; embed_split_idx < params.hidden_split_time; ++embed_split_idx) { - if (embed_split_idx == 4) { - params.embed_split_size = 64; - params.round_embed_split_size = 64; - } - LoadQDataToL0A(params, embed_split_idx, false); - LoadKVDataToL0B(params, embed_split_idx, false); - ComputeQKMMad(params, embed_split_idx, false); - CopyQKResultToGM(params, embed_split_idx, n_idx, false); + } else{ + tail = 0; // control whether to run tail optimization } + uint32_t totalTaskNumRound = totalTaskNum - tail; + + + for (uint32_t process = block_idx; process < totalTaskNumRound; process += (uint32_t)block_num) { // for task + uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; + uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - ComputeQRope(params, n_idx); - FftsCrossCoreSync(QK_READY_DECODER); - } - - // ==================== ComputePV 三层架构拆分 ==================== - // === 第一层：PVParams 结构体 === - struct PVParams { - uint32_t qk_n_2; - uint32_t qk_round_n_2; - uint32_t qk_round_n_2_l1; - uint64_t k_round_n; - uint32_t row_num; - uint64_t hidden_size; - uint32_t l1_kv_pingpong_flag; - uint32_t l0_p_pingpong_flag; - uint32_t embed_split_size; - uint32_t round_embed_split_size; - uint32_t l0c_pingpong_flag; - uint32_t l0b_pingpong_flag; - uint64_t l1kv_offset; - }; - - // ==================== InnerRunCubeMLATP1 三层架构拆分 ==================== - // === 第一层：TP1Context 结构体 === - struct TP1Context { - // input params - uint32_t cur_batch; - uint32_t start_head; - uint32_t cur_head_num; - uint32_t start_kv; - uint32_t cur_q_seqlen; - uint32_t cur_kv_seqlen; - uint32_t offset_tiling; - - // Q address - uint64_t q_offset; - uint64_t q_rope_offset; - - // loop & size - uint32_t pp_n_scalar; - uint32_t sub_n_loop; - uint32_t n_loop; - - // QK dims - uint32_t qk_n; - uint32_t qk_round_n; - uint32_t qk_n_2; - uint32_t qk_round_n_2; - - // row info - uint32_t row_num; - // m is class member - - // TP1 专有字段 - uint32_t sv_n; // PV阶段有效列数（动态变化） - uint32_t s_block_stack; // 流水线深度=4 - }; - - // === 第一层：InitTP1Context 编排层 === - __aicore__ __attribute__((always_inline)) inline void InitTP1Context( - TP1Context &ctx, uint32_t cur_batch, uint32_t start_head, uint32_t cur_head_num, - uint32_t start_kv, uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t offset_tiling) - { - ctx.cur_batch = cur_batch; - ctx.start_head = start_head; - ctx.cur_head_num = cur_head_num; - ctx.start_kv = start_kv; - ctx.cur_q_seqlen = cur_q_seqlen; - ctx.cur_kv_seqlen = cur_kv_seqlen; - ctx.offset_tiling = offset_tiling; - - uint32_t prev_task = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); - uint64_t addr_q_scalar = (uint64_t)prev_task * q_heads; - ctx.q_offset = addr_q_scalar * 512 + start_head * 512; - ctx.q_rope_offset = addr_q_scalar * 64 + start_head * 64; - - ctx.pp_n_scalar = block_size; - ctx.sub_n_loop = ctx.pp_n_scalar / block_size; - ctx.n_loop = (cur_kv_seqlen + ctx.pp_n_scalar - 1) / ctx.pp_n_scalar; - - ctx.qk_n = ctx.pp_n_scalar; - ctx.qk_round_n = RoundUp(ctx.qk_n); - ctx.qk_n_2 = ctx.pp_n_scalar; - ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); - - ctx.row_num = cur_head_num * cur_q_seqlen; - ctx.sv_n = ctx.n_loop == 1 ? cur_kv_seqlen : ctx.pp_n_scalar; - m = RoundUp<16>(ctx.row_num); - - ctx.s_block_stack = 4; - } - - // === 第一层：LoadTP1QData 编排层 === - __aicore__ __attribute__((always_inline)) inline void LoadTP1QData(TP1Context &ctx) - { - LoadQMainFromGMToL1(l1q_buf_addr_tensor, q_gm_tensor, ctx.q_offset, ctx.cur_q_seqlen, ctx.cur_head_num); - LoadQRopeFromGMToL1(l1q_buf_addr_tensor, l1q_rope_buf_addr_tensor, q_rope_gm_tensor, ctx.q_rope_offset, ctx.cur_q_seqlen, ctx.cur_head_num); - PlatformSetQLoadComplete(); - } + uint32_t q_seqlen = 1; + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); + if (kv_seqlen == 0) { + continue; + } + uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; + uint32_t start_head = 0; + uint32_t start_kv = 0; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen; + uint32_t cur_head_num = q_heads; - // [PV 第三层平台函数 + 第二层业务函数] 已迁移至 arch32.h / bs.h - // === 以下 include 第三层和第二层子文件（放在 QKParams/PVParams/TP1Context 之后，确保类型可见）=== - #include "multi_latent_attention_arch32.h" - #include "multi_latent_attention_bs.h" + InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); + } - // === 第一层：ComputeTP1QK 编排层（CUBE1 stage1）=== - // 外层 split_idx 循环 + 内层 embed_split_idx 循环（5次: 0-4） - // idx 0,2: Q→L0A + KV Main GM→L1 - // idx 1,3: Q→L0A + KV L1→L0B + mmad - // idx 4: Q→L0A + KV Rope GM→L1 + L1→L0B + mmad + S→GM - // 循环后: FftsCrossCoreSync - __aicore__ __attribute__((always_inline)) inline void ComputeTP1QK( - TP1Context &ctx, uint32_t n_idx) - { - for (uint32_t split_idx = 0; split_idx < ctx.s_block_stack && n_idx + split_idx < ctx.n_loop; split_idx++) { - TP1QKParams params; - InitTP1QKParams(ctx, split_idx, n_idx, params); - - for (uint32_t embed_split_idx = 0; embed_split_idx < 5; ++embed_split_idx) { - if (embed_split_idx == 4) { - params.embed_split_size = 64; - params.round_embed_split_size = 64; + // suppose all seqs have same length + if (tail > 0){ + uint32_t sample_kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + tiling_head_size + 2)); + bool enableExtraOptimization = true; + if (block_num % 4 == 3) { + // cannot optimize this situation due to math problem + enableExtraOptimization = false; + } + if (sample_kv_seqlen <= 2048){ + // Too Short to benefit from optimization + enableExtraOptimization = false; + } + if (!enableExtraOptimization || tail <= block_num / 2) { + // collect all metadata + uint32_t cores_per_seq = 1; + if (0 < tail && tail <= block_num / 4) {// 6 tasks left, each works with 4 cores + cores_per_seq = 4; + if (tail == 1){ + cores_per_seq = block_num; + } + else if (tail == 2){ + cores_per_seq = block_num / 2; + } + else if(tail == 3){ + cores_per_seq = block_num / 3; + } + else if(tail == 4){ + cores_per_seq = block_num / 4; + } } - - // Q L1→L0A - LoadTP1QDataToL0A(params, embed_split_idx); - - // KV Main GM→L1 (idx 0,2) - if (embed_split_idx == 0 || embed_split_idx == 2) { - LoadTP1KVMainFromGM(params, embed_split_idx); + else if(block_num / 4 < tail && tail <= block_num / 3) { // 8 tasks left, each works with 3 cores + cores_per_seq = 3; + } - - // KV Rope GM→L1 (idx 4) - if (embed_split_idx == 4) { - LoadTP1KVRopeFromGM(params, embed_split_idx); + else if(block_num / 3 < tail && tail <= block_num / 2) { // 12 tasks left, each works with 2 cores + cores_per_seq = 2; + } + else { + // no extra optimization for tail > 12 + cores_per_seq = 1; } - // KV L1→L0B (all idx) - LoadTP1KVDataToL0B(params, embed_split_idx); - - // QK mmad - ComputeTP1QKMMad(params, embed_split_idx); - - // S→GM (idx 4) - if (embed_split_idx == 4) { - CopyTP1QKResultToGM(params, embed_split_idx, n_idx); + if(!enableExtraOptimization){ + cores_per_seq = 1; + } + + uint32_t process = totalTaskNumRound + block_idx / cores_per_seq; + if (process < totalTaskNum) { + uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; + uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + + uint32_t q_seqlen = 1; + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); + uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; + uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; + uint32_t actual_work_cores = kv_seqlen / kv_seqlen_align + (kv_seqlen % kv_seqlen_align != 0); + // cores_per_seq = actual_work_cores; + uint32_t kv_seqlen_process = 0; + if (block_idx < block_idx / cores_per_seq * cores_per_seq + actual_work_cores){ + kv_seqlen_process = (block_idx % cores_per_seq == actual_work_cores - 1) ? + (kv_seqlen - kv_seqlen_align * (actual_work_cores - 1)) : kv_seqlen_align; + } + + if (kv_seqlen > 0 && kv_seqlen_process > 0) { + uint32_t start_head = 0; + uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen_process; + uint32_t cur_head_num = q_heads; + + // no need to modify anything in cube kernel, just call the same kernel + InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); + } } } - } - FftsCrossCoreSync(QK_READY_DECODER); - } - - // === 第一层：ComputePV 编排层 === - __aicore__ __attribute__((always_inline)) inline void ComputePV( - MLAContext &ctx, uint32_t n_idx) - { - PVParams params; - InitPVParams(ctx, n_idx, params); - - for (uint32_t embed_split_idx = 0; embed_split_idx < embed_split_loop_v; ++embed_split_idx) { - params.l0c_pingpong_flag = (n_idx + embed_split_idx) % 2; - params.l0b_pingpong_flag = (embed_split_idx + 1) % 2; - params.l1kv_offset = embed_split_idx * params.k_round_n * params.round_embed_split_size; - - LoadKVTransposeToL0B(params, embed_split_idx); + else if (tail > 3 * block_num / 4){ + // no benefit for optimizing this situation + uint32_t process = totalTaskNumRound + block_idx; + if (process < totalTaskNum) { + uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; + uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - if (embed_split_idx == 0) { - LoadPDataToL0A(params, n_idx); + uint32_t q_seqlen = 1; + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); + if (kv_seqlen > 0) { + uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; + + uint32_t start_head = 0; + uint32_t start_kv = 0; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen; + uint32_t cur_head_num = q_heads; + + InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); + } + } + } + else { + // 18 >= tail >= 12 + // first 12 tasks, two cores per task + { + uint32_t cores_per_seq = 2; + uint32_t process = totalTaskNumRound + block_idx / cores_per_seq; + uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; + uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + + uint32_t q_seqlen = 1; + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); + uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; + uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; + uint32_t kv_seqlen_process = (block_idx % cores_per_seq == cores_per_seq - 1) ? + (kv_seqlen - kv_seqlen_align * (cores_per_seq - 1)) : kv_seqlen_align; + + if (kv_seqlen > 0 && kv_seqlen_process > 0) { + uint32_t start_head = 0; + uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen_process; + uint32_t cur_head_num = q_heads; + + // no need to modify anything in cube kernel, just call the same kernel + InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); + } + } + { + uint32_t cores_per_seq = 4; + uint32_t process = totalTaskNumRound + block_num / 2 + block_idx / cores_per_seq; + if (process < totalTaskNum) { + uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; + uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + + uint32_t q_seqlen = 1; + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); + uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; + uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; + uint32_t kv_seqlen_process = (block_idx % cores_per_seq == cores_per_seq - 1) ? + (kv_seqlen - kv_seqlen_align * (cores_per_seq - 1)) : kv_seqlen_align; + + if (kv_seqlen > 0 && kv_seqlen_process > 0) { + uint32_t start_head = 0; + uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen_process; + uint32_t cur_head_num = q_heads; + + // no need to modify anything in cube kernel, just call the same kernel + InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); + } + } + } } - - ComputePVMmad(params, embed_split_idx); - - CopyPVResultToGM(params, embed_split_idx, n_idx); } - FftsCrossCoreSync(UPDATE_READY_DECODER); + + WAIT_FLAG(M, MTE1, EVENT_ID0); + WAIT_FLAG(M, MTE1, EVENT_ID1); + WAIT_FLAG(M, MTE1, EVENT_ID2); + WAIT_FLAG(M, MTE1, EVENT_ID3); + WAIT_FLAG(M, MTE1, EVENT_ID4); + WAIT_FLAG(M, MTE1, EVENT_ID5); + WAIT_FLAG(M, MTE1, EVENT_ID6); + WAIT_FLAG(M, MTE1, EVENT_ID7); + WAIT_FLAG(FIX, M, EVENT_ID0); + WAIT_FLAG(FIX, M, EVENT_ID1); + WAIT_FLAG(MTE1, MTE2, EVENT_ID0); + WAIT_FLAG(MTE1, MTE2, EVENT_ID1); + WAIT_FLAG(MTE1, MTE2, EVENT_ID2); + WAIT_FLAG(MTE1, MTE2, EVENT_ID3); + WAIT_FLAG(MTE1, MTE2, EVENT_ID4); + WAIT_FLAG(MTE1, MTE2, EVENT_ID5); + WAIT_FLAG(MTE1, MTE2, EVENT_ID6); + WAIT_FLAG(MTE1, MTE2, EVENT_ID7); + WAIT_FLAG(FIX, MTE1, EVENT_ID0); + WAIT_FLAG(FIX, MTE1, EVENT_ID1); + WAIT_FLAG(FIX, MTE1, EVENT_ID2); + WAIT_FLAG(FIX, MTE1, EVENT_ID3); + WAIT_FLAG(FIX, MTE1, EVENT_ID4); + WAIT_FLAG(FIX, MTE1, EVENT_ID5); + WAIT_FLAG(MTE2, FIX, EVENT_ID0); + PIPE_BARRIER(ALL); } +private: __aicore__ __attribute__((always_inline)) inline void InnerRunCubeMLA(uint32_t cur_batch, uint32_t start_head, uint32_t cur_head_num, uint32_t start_kv, uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t offset_tiling) { - MLAContext ctx; - InitMLAContext(ctx, cur_batch, start_head, cur_head_num, - start_kv, cur_q_seqlen, cur_kv_seqlen, offset_tiling); + uint32_t addr_q_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); + uint32_t addr_q_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 3 + offset_tiling)); + uint64_t addr_q_scalar = (uint64_t)(((uint64_t)addr_q_high32) << 32 | addr_q_loww32); + uint64_t q_offset = addr_q_scalar * 512 + start_head * 512; + uint64_t q_rope_offset = addr_q_scalar * 64 + start_head * 64; + + uint32_t pp_n_scalar = block_size; + uint32_t sub_n_loop = pp_n_scalar / block_size; + + uint32_t n_loop = (cur_kv_seqlen + pp_n_scalar - 1) / pp_n_scalar; + + uint32_t qk_n = pp_n_scalar; + uint32_t qk_round_n = RoundUp(qk_n); + uint32_t qk_n_2 = pp_n_scalar; + uint32_t qk_round_n_2 = RoundUp(qk_n_2); + uint32_t qk_round_n_l1 = RoundUp(qk_n); + uint32_t qk_round_n_2_l1 = RoundUp(qk_n_2); + uint64_t hidden_size = 576; + if constexpr(tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + hidden_size = 512; + } + uint64_t k_round_n = qk_round_n; + uint32_t row_num = cur_head_num * cur_q_seqlen; + m = RoundUp<16>(row_num); + + // copy Q + if (cur_q_seqlen == 1) { + gm_to_l1( + l1q_buf_addr_tensor, + q_gm_tensor[q_offset], + cur_head_num, // nValue + RoundUp<16>(cur_head_num),// dstNzC0Stride + 0, // dstNzMatrixStride, unused + 512, // dValue + 0, // dstNzMatrixStride, unused + 512 // srcDValue + ); + } else { + if (q_heads < 128) { + AscendC::DataCopy( + l1q_buf_addr_tensor, + q_gm_tensor[q_offset], + AscendC::Nd2NzParams( + cur_q_seqlen, // ndNum + cur_head_num, // nValue + 512, // dValue + 512 * q_heads, // srcNdMatrixStride + 512, // srcDValue + RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride + cur_q_seqlen, // dstNzNStride + 16 // dstNzMatrixStride + ) + ); + } else { + for (uint32_t ii =0; ii < cur_q_seqlen; ii++) { + AscendC::DataCopy( + l1q_buf_addr_tensor[ii * 16], // offset one datablock + q_gm_tensor[q_offset + ii * q_heads * 512], + AscendC::Nd2NzParams( + 1, // ndNum + cur_head_num, // nValue + 512, // dValue + 0, // srcNdMatrixStride + 512, // srcDValue + RoundUp<16>(cur_q_seqlen * cur_head_num), // dstNzC0Stride + cur_q_seqlen, // dstNzNStride + 16 // dstNzMatrixStride + ) + ); + } + } - LoadQData(ctx); - for (uint32_t n_idx = 0; n_idx < ctx.n_loop + 1; n_idx += 1) { - if (n_idx != ctx.n_loop) { + } + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + gm_to_l1( + l1q_rope_buf_addr_tensor, + q_rope_gm_tensor[q_rope_offset], + cur_head_num, // nValue + RoundUp<16>(cur_head_num),// dstNzC0Stride + 0, // dstNzMatrixStride, unused + 64, // dValue + 0, // dstNzMatrixStride, unused + 64 // srcDValue + ); + } else { + AscendC::DataCopy( + l1q_buf_addr_tensor[RoundUp<16>(cur_head_num * cur_q_seqlen) * 512], + q_rope_gm_tensor[q_rope_offset], + AscendC::Nd2NzParams( + cur_head_num, // ndNum, 32 + cur_q_seqlen, // nValue, 4 + 64, // dValue + 64, // srcNdMatrixStride + 64 * q_heads, // srcDValue + RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride + 1, // dstNzNStride + 16 * cur_q_seqlen // dstNzMatrixStride + ) + ); + } + SET_FLAG(MTE2, MTE1, EVENT_ID0); + WAIT_FLAG(MTE2, MTE1, EVENT_ID0); + for (uint32_t n_idx = 0; n_idx < n_loop + 1; n_idx+=1) { + if (n_idx != n_loop) { uint32_t l1_kv_pingpong_flag = n_idx % 2; - if (n_idx == (ctx.n_loop - 1)) { - ctx.qk_n = (ctx.cur_kv_seqlen - n_idx * ctx.pp_n_scalar); - ctx.qk_round_n = RoundUp(ctx.qk_n); - ctx.qk_round_n_l1 = RoundUp(ctx.qk_n); + if (n_idx == (n_loop - 1)) { + qk_n = (cur_kv_seqlen - n_idx * pp_n_scalar); + qk_round_n = RoundUp(qk_n); + qk_round_n_l1 = RoundUp(qk_n); } - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - ctx.k_round_n = ctx.qk_round_n_l1; + if constexpr(tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + k_round_n = qk_round_n_l1; } else { - ctx.k_round_n = ctx.qk_round_n; + k_round_n = qk_round_n; } + uint64_t hiddenSize_offset = start_head * cur_q_seqlen * embedding_size; + uint32_t embed_split_size = 128; + uint32_t round_embed_split_size = RoundUp(embed_split_size); - LoadKVData(ctx, n_idx, l1_kv_pingpong_flag); - ComputeQK(ctx, n_idx, l1_kv_pingpong_flag); - } - if (n_idx != 0) { - ComputePV(ctx, n_idx); - } - } - } + /* ************ CUBE1 stage1 ************* */ + + uint32_t block_table_id = (uint32_t)(*(block_tables_gm + + cur_batch * max_num_blocks_per_query + start_kv / block_size + n_idx)); + int64_t kv_offset = (int64_t)block_table_id * block_size * stride_kv; + int64_t kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; + uint32_t q_load_coeff = 1; + q_load_coeff = m; + WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); // wait for v -> L0B + if constexpr(KInputType == InputFormat::ND_FORMAT) { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576], + k_gm_tensor[kv_offset], + qk_n, // nValue + qk_round_n, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 512, // dValue + 0, // dstNzMatrixStride, unused + stride_kv // srcDValue + ); + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576 + 512 * qk_round_n], + k_rope_gm_tensor[kv_offset_rope], + qk_n, // nValue + qk_round_n, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 64, // dValue + 0, // dstNzMatrixStride, unused + stride_kv_rope // srcDValue + ); + } else if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 512], + k_gm_tensor[kv_offset], + qk_round_n_l1, // nValue + block_size, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 512, // dValue + 0, // dstNzMatrixStride, unused + 0 // srcDValue + ); + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B + gm_to_l1( + l1kv_rope_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64], + k_rope_gm_tensor[kv_offset_rope], + qk_round_n, // nValue + block_size, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 64, // dValue + 0, // dstNzMatrixStride, unused + 0 // srcDValue + ); + } else { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576], + k_gm_tensor[kv_offset], + qk_round_n, // nValue + block_size, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 512, // dValue + 0, // dstNzMatrixStride, unused + 0 // srcDValue + ); + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); // wait for v -> L0B + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 576 + 512 * qk_round_n], + k_rope_gm_tensor[kv_offset_rope], + qk_round_n, // nValue + block_size, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 64, // dValue + 0, // dstNzMatrixStride, unused + 0 // srcDValue + ); + } + + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); + uint64_t hidden_split_time = (hidden_size + 128 - 1) / 128; + uint64_t embed_split_idx = 0; + for (embed_split_idx = 0; embed_split_idx < hidden_split_time; ++embed_split_idx) { + if (embed_split_idx == 4) { + embed_split_size = 64; + round_embed_split_size = 64; + } + WAIT_FLAG(M, MTE1, embed_split_idx % 2); + +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) + // 3510: L0A is NZ fractal. Single LoadData call with mStep=M/16 + // handles all M-direction fractals. No manual offset loop needed. + l1_to_l0_a( + l0a_buf_tensor[embed_split_idx % 2 * 16384], + l1q_buf_addr_tensor[embed_split_idx * m * 128], + 0, + round_embed_split_size / T_BLOCK_SIZE, // repeat (K-direction) + 0, + q_load_coeff / BLOCK_SIZE, // srcStride (M-direction step) + 0, + 0 // dstStride + ); +#else + for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { + l1_to_l0_a( + l0a_buf_tensor[embed_split_idx % 2 * 16384 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], + l1q_buf_addr_tensor[embed_split_idx * m * 128 + loa_load_idx * T_CUBE_MATRIX_SIZE], + 0, + round_embed_split_size / T_BLOCK_SIZE, // repeat + 0, + q_load_coeff / BLOCK_SIZE, // srcStride + 0, + 0 // dstStride + ); + } +#endif + + SET_FLAG(MTE1, M, embed_split_idx % 2); + + if (embed_split_idx == 0) { + WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + } + if (embed_split_idx == 4) { + WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); + } + WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); + l1_to_l0_b( + l0b_buf_tensor[embed_split_idx % 2 * 16384], + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + embed_split_idx * k_round_n * 128], + 0, + round_embed_split_size * k_round_n / T_CUBE_MATRIX_SIZE, // repeat + 0, + 1, // srcStride + 0, + 0 // dstStride + ); + if (embed_split_idx == 4) { + SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); + } + SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); + WAIT_FLAG(MTE1, M, embed_split_idx % 2); + WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); + if (embed_split_idx == 0) { + WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); + } + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + mmad( + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + l0a_buf_tensor[embed_split_idx % 2 * 16384], + l0b_buf_tensor[embed_split_idx % 2 * 16384], + m, // m + qk_round_n_l1, // n + embed_split_size, // k + embed_split_idx == 0 // cmatrixInitVal + ); + } else { + mmad( + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + l0a_buf_tensor[embed_split_idx % 2 * 16384], + l0b_buf_tensor[embed_split_idx % 2 * 16384], + m, // m + qk_n, // n + embed_split_size, // k + embed_split_idx == 0 // cmatrixInitVal + ); + } + + PIPE_BARRIER(M); + SET_FLAG(M, MTE1, embed_split_idx % 2); + SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); + + // copy S to gm + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + if (embed_split_idx == 3) { + SET_FLAG(M, FIX, l1_kv_pingpong_flag); + WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); + + l0c_to_gm( + s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + m, // MSize + qk_round_n, // NSize + RoundUp<16>(m), // srcStride + qk_round_n // dstStride_dst_D + ); + SET_FLAG(FIX, M, l1_kv_pingpong_flag); + } + } + if (embed_split_idx == 4) { + SET_FLAG(M, FIX, l1_kv_pingpong_flag); + WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); + + // DUMP: dump QK score from L0C before Fixpipe + // AscendC::DumpTensor(mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], 1001, m * qk_round_n); + + l0c_to_gm( + s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + m, // MSize + qk_round_n, // NSize + RoundUp<16>(m), // srcStride + qk_round_n // dstStride_dst_D + ); + SET_FLAG(FIX, M, l1_kv_pingpong_flag); + } + } + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + embed_split_idx = 4; + embed_split_size = 64; + round_embed_split_size = 64; + WAIT_FLAG(M, MTE1, embed_split_idx % 2); + +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) + l1_to_l0_a( + l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], + l1q_rope_buf_addr_tensor[0], + 0, + round_embed_split_size / BLOCK_SIZE, // repeat + 0, + q_load_coeff / BLOCK_SIZE, // srcStride + 0, + 0 // dstStride + ); +#else + for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { + l1_to_l0_a( + l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], + l1q_rope_buf_addr_tensor[loa_load_idx * CUBE_MATRIX_SIZE], + 0, + round_embed_split_size / BLOCK_SIZE, // repeat + 0, + q_load_coeff / BLOCK_SIZE, // srcStride + 0, + 0 // dstStride + ); + } +#endif + + SET_FLAG(MTE1, M, embed_split_idx % 2); + + WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); + WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); + l1_to_l0_b( + l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], + l1kv_rope_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64], + 0, + round_embed_split_size * qk_round_n / CUBE_MATRIX_SIZE, // repeat + 0, + 1, // srcStride + 0, + 0 // dstStride + ); + + SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); + SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); + WAIT_FLAG(MTE1, M, embed_split_idx % 2); + WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); + } + mmad( + mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], + l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], + l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], + m, // m + qk_n, // n + embed_split_size, // k + 1 // cmatrixInitVal + ); + PIPE_BARRIER(M); + SET_FLAG(M, MTE1, embed_split_idx % 2); + SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); + + SET_FLAG(M, FIX, l1_kv_pingpong_flag); + WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + l0c_to_gm( + s_rope_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], + m, // MSize + qk_round_n, // NSize + RoundUp<16>(m), // srcStride + qk_round_n // dstStride_dst_D + ); + } else { + l0c_to_gm( + s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + m, // MSize + qk_round_n, // NSize + RoundUp<16>(m), // srcStride + qk_round_n // dstStride_dst_D + ); + } + SET_FLAG(FIX, M, l1_kv_pingpong_flag); + } + FftsCrossCoreSync(QK_READY_DECODER); + } + /* ************ CUBE2 stage1 ************* */ + if (n_idx != 0) { + if (n_idx == n_loop) { + qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); + qk_round_n_2 = RoundUp(qk_n_2); + qk_round_n_2_l1 = RoundUp(qk_n_2); + } + k_round_n = qk_round_n_2_l1; + uint32_t l1_kv_pingpong_flag = (n_idx - 1) % 2; + uint32_t l0_p_pingpong_flag = (n_idx - 1) % 2; + uint32_t embed_split_size = 128; + embed_split_loop_v = 4; + uint32_t round_embed_split_size = RoundUp(embed_split_size); + for (uint32_t embed_split_idx = 0; embed_split_idx < embed_split_loop_v; ++embed_split_idx) { + uint32_t l0c_pingpong_flag = (n_idx + embed_split_idx) % 2; + uint32_t l0b_pingpong_flag = (embed_split_idx + 1) % 2; + uint64_t l1kv_offset = embed_split_idx * k_round_n * round_embed_split_size; + WAIT_FLAG(M, MTE1, l0b_pingpong_flag + 2); + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.dstGap = 0; + loadDataParams.startIndex = 0; + loadDataParams.dstFracGap = 0; + if (k_round_n <= round_embed_split_size) { // Nz -> nZ + loadDataParams.repeatTimes = round_embed_split_size / T_BLOCK_SIZE; + loadDataParams.srcStride = k_round_n / T_BLOCK_SIZE; + uint16_t dstGap = sizeof(IN_DTYPE) == 1 ? 1 : 0; + loadDataParams.dstGap = dstGap; + for (uint32_t l0b_load_idx = 0; l0b_load_idx < k_round_n / T_BLOCK_SIZE; ++l0b_load_idx) { + // along embd dim + AscendC::LoadDataWithTranspose( + l0b_buf_tensor[l0b_pingpong_flag * 16384 + l0b_load_idx * RoundUp<16>(embed_split_size) * T_BLOCK_SIZE], + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + l1kv_offset + l0b_load_idx * T_BLOCK_SIZE * T_BLOCK_SIZE], + loadDataParams); + } + } else { + for (uint32_t l0b_load_idx = 0; l0b_load_idx < round_embed_split_size / T_BLOCK_SIZE; ++l0b_load_idx) { + // along kv_len_blk dim + loadDataParams.repeatTimes = qk_round_n_2 / T_BLOCK_SIZE; + loadDataParams.srcStride = 1; + loadDataParams.dstGap = round_embed_split_size / BLOCK_SIZE - 1; + AscendC::LoadDataWithTranspose( + l0b_buf_tensor[l0b_pingpong_flag * 16384 + l0b_load_idx * T_BLOCK_SIZE * T_BLOCK_SIZE], + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + l1kv_offset + l0b_load_idx * qk_round_n_2 * T_BLOCK_SIZE], + loadDataParams); + } + } + if (embed_split_idx == embed_split_loop_v - 1) { + SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); + } + // move p from gm to l1 + uint32_t p_move_head_num = row_num; + if (embed_split_idx == 0) { + WaitFlagDev(SOFTMAX_READY_DECODER); + + WAIT_FLAG(MTE1, MTE2, EVENT_ID7); + gm_to_l1( + l1p_buf_addr_tensor, + p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + ((n_idx - 1) % 2) * TMP_SIZE * T_BLOCK_OFFSET / 2], + p_move_head_num, // nValue + RoundUp(p_move_head_num),// dstNzC0Stride + 0, // dstNzMatrixStride, unused + k_round_n, // dValue + 0, // dstNzMatrixStride, unused + qk_round_n_2 * 2 / sizeof(IN_DTYPE) // srcDValue + ); + SET_FLAG(MTE2, MTE1, EVENT_ID7); + WAIT_FLAG(MTE2, MTE1, EVENT_ID7); + // move p from l1 to l0a + WAIT_FLAG(M, MTE1, l0_p_pingpong_flag); + uint32_t p_load_coeff = RoundUp<16>(p_move_head_num); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + l1_to_l0_a( + l0a_buf_tensor[l0_p_pingpong_flag * 16384], l1p_buf_addr_tensor, RoundUp(p_move_head_num), + qk_round_n_2_l1, // repeat + 0, + 0, // srcStride + 0, + 0 // dstStride + ); + } else { +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) + l1_to_l0_a( + l0a_buf_tensor[l0_p_pingpong_flag * 16384], + l1p_buf_addr_tensor[0], + 0, + qk_round_n_2 / T_BLOCK_SIZE, // repeat + 0, + p_load_coeff / BLOCK_SIZE, // srcStride + 0, + 0 // dstStride + ); +#else + for (uint64_t loa_load_idx = 0; loa_load_idx < p_load_coeff / BLOCK_SIZE; ++loa_load_idx) { + l1_to_l0_a( + l0a_buf_tensor[l0_p_pingpong_flag * 16384 + loa_load_idx * qk_round_n_2 * BLOCK_SIZE], + l1p_buf_addr_tensor[loa_load_idx * T_CUBE_MATRIX_SIZE], + 0, + qk_round_n_2 / T_BLOCK_SIZE, // repeat + 0, + p_load_coeff / BLOCK_SIZE, // srcStride + 0, + 0 // dstStride + ); + } +#endif + } + SET_FLAG(MTE1, MTE2, EVENT_ID7); + } + SET_FLAG(MTE1, M, l0b_pingpong_flag); + WAIT_FLAG(MTE1, M, l0b_pingpong_flag); + WAIT_FLAG(FIX, M, l0c_pingpong_flag); + mmad( + mm2_l0c_buf_tensor[l0c_pingpong_flag * 16384], + l0a_buf_tensor[l0_p_pingpong_flag * 16384], + l0b_buf_tensor[l0b_pingpong_flag * 16384], + m, // m + embed_split_size, // n + qk_n_2, // k + 1 // cmatrixInitVal + ); + SET_FLAG(M, MTE1, l0b_pingpong_flag + 2); + if (embed_split_idx == embed_split_loop_v - 1) { + SET_FLAG(M, MTE1, l0_p_pingpong_flag); + } + SET_FLAG(M, FIX, l0c_pingpong_flag); + WAIT_FLAG(M, FIX, l0c_pingpong_flag); + + // copy O to gm + l0c_to_gm( + o_tmp_gm_tensor[(uint64_t)block_idx * TMP_SIZE * 2 + embed_split_idx * round_embed_split_size + ((n_idx - 1) % 2) * TMP_SIZE], + mm2_l0c_buf_tensor[l0c_pingpong_flag * 16384], + m, // MSize + RoundUp<16>(embed_split_size), // NSize 32B align + RoundUp<16>(m), // srcStride + round_v // dstStride_dst_D + ); + SET_FLAG(FIX, M, l0c_pingpong_flag); + } + FftsCrossCoreSync(UPDATE_READY_DECODER); + } + } + } __aicore__ __attribute__((always_inline)) inline void InnerRunCubeMLATP1(uint32_t cur_batch, uint32_t start_head, uint32_t cur_head_num, uint32_t start_kv, uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t offset_tiling) { - TP1Context ctx; - InitTP1Context(ctx, cur_batch, start_head, cur_head_num, - start_kv, cur_q_seqlen, cur_kv_seqlen, offset_tiling); - LoadTP1QData(ctx); - - uint32_t pp_n_scalar = ctx.pp_n_scalar; - uint32_t n_loop = ctx.n_loop; - uint32_t s_block_stack = ctx.s_block_stack; + uint32_t prev_task = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); + uint64_t addr_q_scalar = (uint64_t)prev_task * q_heads; + uint64_t q_offset = addr_q_scalar * 512 + start_head * 512; + uint64_t q_rope_offset = addr_q_scalar * 64 + start_head * 64; + + uint32_t pp_n_scalar = block_size; + uint32_t sub_n_loop = pp_n_scalar / block_size; + + uint32_t n_loop = (cur_kv_seqlen + pp_n_scalar - 1) / pp_n_scalar; + + uint32_t qk_n = pp_n_scalar; + uint32_t qk_round_n = RoundUp(qk_n); + uint32_t qk_n_2 = pp_n_scalar; + uint32_t qk_round_n_2 = RoundUp(qk_n_2); + + uint32_t row_num = cur_head_num * cur_q_seqlen; + + uint32_t sv_n = n_loop == 1 ? cur_kv_seqlen : pp_n_scalar; + m = RoundUp<16>(row_num); + + // copy Q + if (cur_q_seqlen == 1) { + gm_to_l1( + l1q_buf_addr_tensor, + q_gm_tensor[q_offset], + cur_head_num, // nValue + RoundUp<16>(cur_head_num),// dstNzC0Stride + 0, // dstNzMatrixStride, unused + 512, // dValue + 0, // dstNzMatrixStride, unused + 512 // srcDValue + ); + } else { + if (q_heads < 128) { + AscendC::DataCopy( + l1q_buf_addr_tensor, + q_gm_tensor[q_offset], + AscendC::Nd2NzParams( + cur_q_seqlen, // ndNum + cur_head_num, // nValue + 512, // dValue + 512 * q_heads, // srcNdMatrixStride + 512, // srcDValue + RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride + cur_q_seqlen, // dstNzNStride + 16 // dstNzMatrixStride + ) + ); + } else { + for (uint32_t ii =0; ii < cur_q_seqlen; ii++) { + AscendC::DataCopy( + l1q_buf_addr_tensor[ii * 16], // offset one datablock + q_gm_tensor[q_offset + ii * q_heads * 512], + AscendC::Nd2NzParams( + 1, // ndNum + cur_head_num, // nValue + 512, // dValue + 0, // srcNdMatrixStride + 512, // srcDValue + RoundUp<16>(cur_q_seqlen * cur_head_num), // dstNzC0Stride + cur_q_seqlen, // dstNzNStride + 16 // dstNzMatrixStride + ) + ); + } + } + + } + + AscendC::DataCopy( + l1q_buf_addr_tensor[RoundUp<16>(cur_head_num * cur_q_seqlen) * 512], + q_rope_gm_tensor[q_rope_offset], + AscendC::Nd2NzParams( + cur_head_num, // ndNum, 32 + cur_q_seqlen, // nValue, 4 + 64, // dValue + 64, // srcNdMatrixStride + 64 * q_heads, // srcDValue + RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride + 1, // dstNzNStride + 16 * cur_q_seqlen // dstNzMatrixStride + ) + ); + + SET_FLAG(MTE2, MTE1, EVENT_ID0); + WAIT_FLAG(MTE2, MTE1, EVENT_ID0); + uint32_t s_block_stack = 4; for (uint32_t n_idx = 0; n_idx < n_loop + s_block_stack; n_idx+=s_block_stack) { if (n_idx < n_loop) { - ComputeTP1QK(ctx, n_idx); + uint32_t sv_n_triu = n_loop * pp_n_scalar; + if (n_idx + s_block_stack > n_loop - 1) { + sv_n = cur_kv_seqlen - n_idx * pp_n_scalar; // delete + } else { + sv_n = pp_n_scalar * s_block_stack; + } + uint32_t sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; + for (uint32_t split_idx = 0; split_idx < s_block_stack && n_idx + split_idx < n_loop; split_idx++) { + uint32_t now_idx = n_idx + split_idx; + uint32_t l1_kv_pingpong_flag = now_idx % 2; + if (now_idx == (n_loop - 1)) { + qk_n = (cur_kv_seqlen - now_idx * pp_n_scalar); + qk_round_n = RoundUp(qk_n); + } else { + qk_n = pp_n_scalar; + qk_round_n = RoundUp(qk_n); + } + bool last_split = split_idx == s_block_stack - 1 || now_idx == n_loop - 1; + uint32_t embed_split_size = 128; + uint32_t round_embed_split_size = RoundUp(embed_split_size); + + /* ************ CUBE1 stage1 ************* */ + + uint32_t block_table_id = (uint32_t)(*(block_tables_gm + + cur_batch * max_num_blocks_per_query + start_kv / block_size + now_idx)); + int64_t kv_offset = (int64_t)block_table_id * block_size * stride_kv; + int64_t kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; + + + uint32_t q_load_coeff = 1; + q_load_coeff = m; + int64_t now_l1_offset = 0; + for (uint32_t embed_split_idx = 0; embed_split_idx < 5; ++embed_split_idx) { + if (embed_split_idx == 4) { + embed_split_size = 64; + round_embed_split_size = 64; + } + WAIT_FLAG(M, MTE1, embed_split_idx % 2); + +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) + l1_to_l0_a( + l0a_buf_tensor[embed_split_idx % 2 * 16384], + l1q_buf_addr_tensor[embed_split_idx * m * 128], + 0, + round_embed_split_size / T_BLOCK_SIZE, // repeat + 0, + q_load_coeff / BLOCK_SIZE, // srcStride + 0, + 0 // dstStride + ); +#else + for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { + l1_to_l0_a( + l0a_buf_tensor[embed_split_idx % 2 * 16384 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], + l1q_buf_addr_tensor[embed_split_idx * m * 128 + loa_load_idx * T_CUBE_MATRIX_SIZE], + 0, + round_embed_split_size / T_BLOCK_SIZE, // repeat + 0, + q_load_coeff / BLOCK_SIZE, // srcStride + 0, + 0 // dstStride + ); + } +#endif + + SET_FLAG(MTE1, M, embed_split_idx % 2); + if (embed_split_idx == 0 || embed_split_idx == 2) { + WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); // 等待V全部搬入L0B + now_l1_offset = l1_kv_pingpong_flag * 128 * 256; + if constexpr(KInputType == InputFormat::ND_FORMAT) { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 256], + k_gm_tensor[kv_offset + embed_split_idx * 128], + qk_n, // nValue + qk_round_n, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 256, // dValue + 0, // dstNzMatrixStride, unused + stride_kv // srcDValue + ); + } else { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 256], + k_gm_tensor[kv_offset + block_size * 128 * embed_split_idx], + qk_round_n, + block_size, + qk_round_n, + 256, + 256, 256); + } + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + } else if (embed_split_idx == 4) { + WAIT_FLAG(MTE1, MTE2, 2 + l1_kv_pingpong_flag); + now_l1_offset = l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; + if constexpr(KInputType == InputFormat::ND_FORMAT) { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128], + k_rope_gm_tensor[kv_offset_rope], + qk_n, // nValue + qk_round_n, // dstNzC0Stride + 0, // dstNzMatrixStride, unused + 64, // dValue + 0, // dstNzMatrixStride, unused + stride_kv_rope // srcDValue + ); + } else { + gm_to_l1( + l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128], + k_rope_gm_tensor[kv_offset_rope], + qk_round_n, + block_size, + qk_round_n, + 64, + 64, + 64); + } + SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); + } + WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); + l1_to_l0_b( + l0b_buf_tensor[embed_split_idx % 2 * 16384], + l1kv_buf_addr_tensor[now_l1_offset + embed_split_idx % 2 * qk_round_n * 128], + 0, + round_embed_split_size * qk_round_n / T_CUBE_MATRIX_SIZE, // repeat + 0, + 1, // srcStride + 0, + 0 // dstStride + ); + if (embed_split_idx == 1 || embed_split_idx == 3) { + SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); + } + if (embed_split_idx == 4) { + SET_FLAG(MTE1, MTE2, 2 + l1_kv_pingpong_flag); + } + SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); + WAIT_FLAG(MTE1, M, embed_split_idx % 2); + WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); + if (embed_split_idx == 0) { + WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); + } + mmad( + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + l0a_buf_tensor[embed_split_idx % 2 * 16384], + l0b_buf_tensor[embed_split_idx % 2 * 16384], + m, // m + qk_n, // n + embed_split_size, // k + embed_split_idx == 0 // cmatrixInitVal + ); + PIPE_BARRIER(M); + SET_FLAG(M, MTE1, embed_split_idx % 2); + SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); + + // copy S to gm + if (embed_split_idx == 4) { + SET_FLAG(M, FIX, l1_kv_pingpong_flag); + WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); + + l0c_to_gm( + s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER * 4 + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar], + mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], + m, // MSize + qk_round_n, // NSize + RoundUp<16>(m), // srcStride + sv_round_n // dstStride_dst_D + ); + SET_FLAG(FIX, M, l1_kv_pingpong_flag); + } + } + } + FftsCrossCoreSync(QK_READY_DECODER); } /* ************ CUBE2 stage1 ************* */ if (n_idx >= s_block_stack) { if (n_idx + s_block_stack > n_loop + s_block_stack - 1) { - ctx.sv_n = cur_kv_seqlen - (n_idx - s_block_stack) * pp_n_scalar; // delete + sv_n = cur_kv_seqlen - (n_idx - s_block_stack) * pp_n_scalar; // delete } else { - ctx.sv_n = pp_n_scalar * s_block_stack; + sv_n = pp_n_scalar * s_block_stack; } - uint32_t sv_n = ctx.sv_n; - uint32_t row_num = ctx.row_num; uint32_t sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; uint32_t embed_split_size = 128; embed_split_loop_v = 4; @@ -815,13 +1539,11 @@ class MLAttentionDecoderAic { for (uint32_t split_idx = 0; split_idx < s_block_stack && n_idx + split_idx < n_loop + s_block_stack; split_idx++) { uint32_t now_idx = n_idx + split_idx; if (now_idx == (n_loop + s_block_stack - 1)) { - ctx.qk_n_2 = (cur_kv_seqlen - (now_idx - s_block_stack) * pp_n_scalar); + qk_n_2 = (cur_kv_seqlen - (now_idx - s_block_stack) * pp_n_scalar); } else { - ctx.qk_n_2 = pp_n_scalar; + qk_n_2 = pp_n_scalar; } - ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); - uint32_t qk_n_2 = ctx.qk_n_2; - uint32_t qk_round_n_2 = ctx.qk_round_n_2; + qk_round_n_2 = RoundUp(qk_n_2); uint32_t l1_kv_pingpong_flag = now_idx % 2; uint32_t l0_p_pingpong_flag = now_idx % 2; uint32_t l0b_pingpong_flag = now_idx % 2; @@ -892,6 +1614,18 @@ class MLAttentionDecoderAic { // move p from l1 to l0a WAIT_FLAG(M, MTE1, l0_p_pingpong_flag); uint32_t p_load_coeff = RoundUp<16>(p_move_head_num); +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) + l1_to_l0_a( + l0a_buf_tensor[l0_p_pingpong_flag * 16384], + l1p_buf_addr_tensor[l0_p_pingpong_flag * 128 * 128], + 0, + qk_round_n_2 / T_BLOCK_SIZE, // repeat + 0, + p_load_coeff / BLOCK_SIZE, // srcStride + 0, + 0 // dstStride + ); +#else for (uint64_t loa_load_idx = 0; loa_load_idx < p_load_coeff / BLOCK_SIZE; ++loa_load_idx) { l1_to_l0_a( l0a_buf_tensor[l0_p_pingpong_flag * 16384 + loa_load_idx * qk_round_n_2 * BLOCK_SIZE], @@ -904,6 +1638,7 @@ class MLAttentionDecoderAic { 0 // dstStride ); } +#endif SET_FLAG(MTE1, MTE2, l0_p_pingpong_flag + 6); SET_FLAG(MTE1, M, l0b_pingpong_flag); @@ -1127,14 +1862,73 @@ class MLADecoderAiv{ __aicore__ __attribute__((always_inline)) inline void Run() { - PlatformInitVectorPipeSync(); - ScheduleVectorTasks(); - PlatformWaitVectorPipeSync(); + // 3510 guard: AIV block_idx may exceed blockDim (cube count) in MIX_AIC mode. + if ((uint32_t)(int32_t)block_idx >= (uint32_t)(int32_t)block_num) { + return; + } + SET_FLAG(MTE3, V, EVENT_ID0); + SET_FLAG(MTE3, MTE2, EVENT_ID0); + SET_FLAG(MTE3, MTE2, EVENT_ID2); + SET_FLAG(MTE3, MTE2, EVENT_ID3); + SET_FLAG(MTE3, MTE2, EVENT_ID4); + SET_FLAG(V, MTE2, EVENT_ID4); + SET_FLAG(V, MTE2, EVENT_ID0); + SET_FLAG(MTE3, V, EVENT_ID2); + SET_FLAG(V, MTE2, EVENT_ID2); + + + uint64_t cur_batch = 0; + + uint32_t q_block_num_per_batch = (q_heads + cur_qn_blk_size - 1) / cur_qn_blk_size; + uint32_t process_num = q_block_num_per_batch * num_batches; + + for (uint32_t process = block_idx; process < process_num; process += (uint32_t)block_num) { // for task + cur_batch = process / q_block_num_per_batch; + if (cur_batch >= num_batches) break; + + uint32_t offset_tiling = tiling_head_size + tiling_para_size * cur_batch; + uint32_t start_core_idx = (cur_batch * q_block_num_per_batch) % block_num; + + uint32_t q_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); + if (kv_seqlen == 0) { + continue; + } + uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; + + uint32_t start_head = (process % q_block_num_per_batch) * cur_qn_blk_size; + uint32_t start_kv = 0; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen; + uint32_t cur_head_num = cur_qn_blk_size; + uint32_t cur_nIndx = 0; + InnerRunVectorChange(cur_batch, start_head, cur_nIndx, cur_q_seq_len, cur_kv_seqlen, cur_head_num, + offset_tiling, 512, embed_split_loop_v_former); + } + + WAIT_FLAG(MTE3, V, EVENT_ID0); + WAIT_FLAG(MTE3, MTE2, EVENT_ID0); + WAIT_FLAG(MTE3, MTE2, EVENT_ID2); + WAIT_FLAG(MTE3, MTE2, EVENT_ID3); + WAIT_FLAG(MTE3, MTE2, EVENT_ID4); + WAIT_FLAG(V, MTE2, EVENT_ID0); + WAIT_FLAG(V, MTE2, EVENT_ID4); + WAIT_FLAG(MTE3, V, EVENT_ID2); + WAIT_FLAG(V, MTE2, EVENT_ID2); } __aicore__ __attribute__((always_inline)) inline void RunTP1() { - PlatformInitVectorPipeSyncTP1(); + SET_FLAG(MTE3, V, EVENT_ID0); + SET_FLAG(MTE3, MTE2, EVENT_ID0); + SET_FLAG(MTE3, MTE2, EVENT_ID1); + SET_FLAG(MTE3, MTE2, EVENT_ID2); + SET_FLAG(MTE3, MTE2, EVENT_ID3); + SET_FLAG(MTE3, MTE2, EVENT_ID4); + SET_FLAG(V, MTE2, EVENT_ID4); + SET_FLAG(V, MTE2, EVENT_ID0); + SET_FLAG(MTE3, V, EVENT_ID2); + SET_FLAG(V, MTE2, EVENT_ID2); uint32_t tail = totalTaskNum % block_num; if constexpr (EnableOptimization) { @@ -1294,118 +2088,1771 @@ class MLADecoderAiv{ } } } - else { - // 18 >= tail >= 12 - // first 12 tasks, two cores per task - { - uint32_t cores_per_seq = 2; - uint32_t process = totalTaskNumRound + block_idx / cores_per_seq; - uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; - uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - - uint32_t q_seqlen = 1; - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); - uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; - uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; - uint32_t kv_seqlen_process = (block_idx % cores_per_seq == cores_per_seq - 1) ? - (kv_seqlen - kv_seqlen_align * (cores_per_seq - 1)) : kv_seqlen_align; - - if (kv_seqlen > 0 && kv_seqlen_process > 0) { - uint32_t start_head = 0; - uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen_process; - uint32_t cur_head_num = q_heads; - uint32_t cur_nIndx = 0; // no use, follow the previous code - - // customized InnerRunVectorChange for tail processing - TailInnerRunVectorChangeTP1(start_head, cur_q_seq_len, cur_kv_seqlen, cur_head_num, offset_tiling, 512, embed_split_loop_v_former); - // Sync all vector cores - AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); - AscendC::CrossCoreWaitFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); - // 4 vector cores - if (block_idx % cores_per_seq < 2) { - TailInnerGatherVectorTP1(start_head + cur_head_num / 2 * (block_idx % cores_per_seq), cur_q_seq_len, cur_head_num / 2, block_idx - (block_idx % cores_per_seq), cores_per_seq, offset_tiling); - } - } - else{ - AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); - AscendC::CrossCoreWaitFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); - } - } + else { + // 18 >= tail >= 12 + // first 12 tasks, two cores per task + { + uint32_t cores_per_seq = 2; + uint32_t process = totalTaskNumRound + block_idx / cores_per_seq; + uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; + uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + + uint32_t q_seqlen = 1; + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); + uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; + uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; + uint32_t kv_seqlen_process = (block_idx % cores_per_seq == cores_per_seq - 1) ? + (kv_seqlen - kv_seqlen_align * (cores_per_seq - 1)) : kv_seqlen_align; + + if (kv_seqlen > 0 && kv_seqlen_process > 0) { + uint32_t start_head = 0; + uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen_process; + uint32_t cur_head_num = q_heads; + uint32_t cur_nIndx = 0; // no use, follow the previous code + + // customized InnerRunVectorChange for tail processing + TailInnerRunVectorChangeTP1(start_head, cur_q_seq_len, cur_kv_seqlen, cur_head_num, offset_tiling, 512, embed_split_loop_v_former); + // Sync all vector cores + AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); + AscendC::CrossCoreWaitFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); + // 4 vector cores + if (block_idx % cores_per_seq < 2) { + TailInnerGatherVectorTP1(start_head + cur_head_num / 2 * (block_idx % cores_per_seq), cur_q_seq_len, cur_head_num / 2, block_idx - (block_idx % cores_per_seq), cores_per_seq, offset_tiling); + } + } + else{ + AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); + AscendC::CrossCoreWaitFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); + } + } + + // other tasks + { + uint32_t cores_per_seq = 4; + // TODO: check whether need extra sync between aic and aiv for used s_gm + uint32_t process = totalTaskNumRound + block_num / 2 + block_idx / cores_per_seq; + if (process < totalTaskNum) { + uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; + uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); + + uint32_t q_seqlen = 1; + uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); + uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; + uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; + uint32_t kv_seqlen_process = (block_idx % cores_per_seq == cores_per_seq - 1) ? + (kv_seqlen - kv_seqlen_align * (cores_per_seq - 1)) : kv_seqlen_align; + + if (kv_seqlen > 0 && kv_seqlen_process > 0) { + uint32_t start_head = 0; + uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; + uint32_t cur_q_seq_len = q_seqlen; + uint32_t cur_kv_seqlen = kv_seqlen_process; + uint32_t cur_head_num = q_heads; + uint32_t cur_nIndx = 0; // no use, follow the previous code + // customized InnerRunVectorChange for tail processing + TailInnerRunVectorChangeTP1(start_head, cur_q_seq_len, cur_kv_seqlen, cur_head_num, offset_tiling, 512, embed_split_loop_v_former); + // Sync all vector cores + AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); + AscendC::CrossCoreWaitFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); + + // 8 vector cores + if (block_idx % cores_per_seq < 4) { + TailInnerGatherVectorTP1(start_head + cur_head_num / 4 * (block_idx % cores_per_seq), cur_q_seq_len, cur_head_num / 4, block_idx - (block_idx % cores_per_seq), cores_per_seq, offset_tiling); + } + + } + else{ + AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); + AscendC::CrossCoreWaitFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); + } + } + else{ + AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); + AscendC::CrossCoreWaitFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); + } + } + } + } + + WAIT_FLAG(MTE3, V, EVENT_ID0); + WAIT_FLAG(MTE3, MTE2, EVENT_ID0); + WAIT_FLAG(MTE3, MTE2, EVENT_ID1); + WAIT_FLAG(MTE3, MTE2, EVENT_ID2); + WAIT_FLAG(MTE3, MTE2, EVENT_ID3); + WAIT_FLAG(MTE3, MTE2, EVENT_ID4); + WAIT_FLAG(V, MTE2, EVENT_ID0); + WAIT_FLAG(V, MTE2, EVENT_ID4); + WAIT_FLAG(MTE3, V, EVENT_ID2); + WAIT_FLAG(V, MTE2, EVENT_ID2); + } +private: + + + __aicore__ __attribute__((always_inline)) inline void ReduceMaxRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& tempTensor, + uint32_t sub_m, + uint32_t qk_n, + uint32_t qk_round_n) + { + if (qk_n <= FLOAT_VECTOR_SIZE) { + __set_mask(qk_n); + cmax_v(dst, + src, + sub_m, // repeat + 1, // dstRepeatStride + 1, // srcBlockStride + qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride + ); + } else { + ub_to_ub( + tempTensor, + src, + 0, // sid + sub_m, // nBurst + HALF_VECTOR_SIZE / BLOCK_SIZE, // lenBurst + (qk_round_n - FLOAT_VECTOR_SIZE) / FLOAT_BLOCK_SIZE, // srcGap + 0 // dstGap + ); + PIPE_BARRIER(V); + for (uint32_t rowmax_idx = 1; rowmax_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowmax_idx) { + max_v( + tempTensor, + tempTensor, + src[rowmax_idx * FLOAT_VECTOR_SIZE], + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride + ); + PIPE_BARRIER(V); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + max_v( + tempTensor, + tempTensor, + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride + ); + } + PIPE_BARRIER(V); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + cmax_v( + dst, + tempTensor, + sub_m, // repeat + 1, // dstRepeatStride + 1, // srcBlockStride + 8 // srcRepeatStride + ); + } + SetVectorMask((uint64_t)-1, (uint64_t)-1); + PIPE_BARRIER(V); + } + + + __aicore__ __attribute__((always_inline)) inline void ReduceSumRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + uint32_t sub_m, + uint32_t qk_n, + uint32_t qk_round_n) + { + if (qk_n <= FLOAT_VECTOR_SIZE) { + __set_mask(qk_n); + cadd_v( + dst, + src, + sub_m, // repeat + 1, // dstRepeatStride + 1, // srcBlockStride + qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else { + for (uint32_t rowsum_idx = 1; rowsum_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowsum_idx) { + add_v( + src, + src, + src[rowsum_idx * FLOAT_VECTOR_SIZE], + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride + ); + PIPE_BARRIER(V); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + add_v( + src, + src, + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + qk_round_n / FLOAT_BLOCK_SIZE // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + + cadd_v( + dst, + src, + sub_m, // repeat + 1, // dstRepeatStride + 1, // srcBlockStride + qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride + ); + } + } + + __aicore__ __attribute__((always_inline)) inline void TensorSubValueRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& MaxTensor, + const AscendC::LocalTensor& tempMaxTensor, + uint32_t sub_m, + uint32_t round_sub_m, + uint32_t qk_n, + uint32_t qk_round_n) + { + brcb_v( + tempMaxTensor.ReinterpretCast(), + MaxTensor.ReinterpretCast(), + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + for (uint32_t sub_v_idx = 0; sub_v_idx < qk_n / FLOAT_VECTOR_SIZE; ++sub_v_idx) { + sub_v(dst[sub_v_idx * FLOAT_VECTOR_SIZE], + src[sub_v_idx * FLOAT_VECTOR_SIZE], + tempMaxTensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + sub_v(dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tempMaxTensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + } + + __aicore__ __attribute__((always_inline)) inline void TensorDivRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& src1, + uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n) + { + PIPE_BARRIER(V); + for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { + div_v(dst[vadd_idx * FLOAT_VECTOR_SIZE], + src[vadd_idx * FLOAT_VECTOR_SIZE], + src1, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + div_v(dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src1, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + } + + __aicore__ __attribute__((always_inline)) inline void TensorMulRepeatM( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& src1, + uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t src1BlockStride + ) { + PIPE_BARRIER(V); + for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { + mul_v(dst[vadd_idx * FLOAT_VECTOR_SIZE], + src[vadd_idx * FLOAT_VECTOR_SIZE], + src1, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + src1BlockStride, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + mul_v(dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + src1, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + src1BlockStride, // src1BlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + } + + __aicore__ __attribute__((always_inline)) inline void DeQuantPerHeadImpl( + const AscendC::GlobalTensor& deScaleGm, + const AscendC::GlobalTensor& src, + AscendC::LocalTensor dst, + AscendC::LocalTensor temp, + AscendC::LocalTensor deScaleUb, + AscendC::LocalTensor tempScale, + AscendC::LocalTensor quantScale, + uint32_t sub_m, + uint32_t qk_n, + uint32_t qk_round_n, + bool online, + bool move_tensor + ){ + gm_to_ub_align(deScaleUb, + deScaleGm, + 0, // sid + 1, // nBurst + sub_m * sizeof(mmScaleType), // lenBurst + 0, // leftPaddingNum + 0, // rightPaddingNum + 0, // srcGap + 0 // dstGap + ); + if (online) { + // if dequant online need mul p quant scale + SET_FLAG(MTE2, V, EVENT_ID2); + WAIT_FLAG(MTE2, V, EVENT_ID2); + TensorMulRepeatM(deScaleUb, deScaleUb, quantScale, 1, sub_m, RoundUp<16>(sub_m), 1); + } + + if (move_tensor) { + gm_to_ub( + temp, + src, + 0, // sid + 1, // nBurst + CeilDiv(sub_m * qk_round_n), // lenBurst + 0, // srcGap + 0 // dstGap + ); + } + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + brcb_v( + tempScale.template ReinterpretCast(), + deScaleUb.template ReinterpretCast(), + 1, // dstBlockStrides + 8, // dstRepeatStride + RoundUp<16>(sub_m) / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + uint32_t count = sub_m * qk_round_n; + uint32_t repeat_times = (count + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE; + if (repeat_times < 255) { + conv_v( + dst, // dst + temp, // src + repeat_times, // repeat_times + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + } else { + for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { + conv_v( + dst[vconv_idx * count / 2], // dst + temp[vconv_idx * count / 2], // src + (count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat_times + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + } + } + TensorMulRepeatM(dst, dst, tempScale, sub_m, qk_n, qk_round_n, 0); + PIPE_BARRIER(V); + } + + __aicore__ __attribute__((always_inline)) inline void QuantPerTokenImpl( + const AscendC::LocalTensor& dst, + const AscendC::LocalTensor& src, + const AscendC::LocalTensor& scale, + uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t pQuantOnline) + { + if (pQuantOnline) { + // scr / scale + TensorDivRepeatM(dst.template ReinterpretCast(), src, scale, sub_m, qk_n, qk_round_n); + } else { + // scr * scale + TensorMulRepeatM(dst.template ReinterpretCast(), src, scale, sub_m, qk_n, qk_round_n, 0); + } + // src fp32 -> casttofp16 -> casttoint8 + uint32_t count = sub_m * qk_round_n; + uint32_t repeat_times = (count + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE; + if (repeat_times < 255) { + conv_v( + dst.template ReinterpretCast(), // dst + dst.template ReinterpretCast(), // src + repeat_times, // repeat_times + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + } else { + for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { + conv_v( + dst.template ReinterpretCast()[vconv_idx * count / 2], // dst + dst.template ReinterpretCast()[vconv_idx * count / 2], // src + (count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat_times + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + } + } + PIPE_BARRIER(V); + for (uint32_t row_idx = 0; row_idx < qk_n / HALF_VECTOR_SIZE; ++row_idx) { + AscendC::Cast(dst.template ReinterpretCast()[row_idx * HALF_VECTOR_SIZE], + dst.template ReinterpretCast()[row_idx * HALF_VECTOR_SIZE], AscendC::RoundMode::CAST_RINT, + (uint64_t)0, sub_m, {1, 1, (uint8_t)((qk_round_n) / BLOCK_SIZE), (uint8_t)(qk_round_n / BLOCK_SIZE)}); + } + if (qk_n % HALF_VECTOR_SIZE > 0) { + __set_mask(qk_n % HALF_VECTOR_SIZE); + AscendC::Cast(dst.template ReinterpretCast()[qk_n / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], + dst.template ReinterpretCast()[qk_n / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], AscendC::RoundMode::CAST_RINT, + (uint64_t)0, sub_m, {1, 1, (uint8_t)((qk_round_n) / BLOCK_SIZE), (uint8_t)(qk_round_n / BLOCK_SIZE)}); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + } + + __aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( + AscendC::GlobalTensor p_gm_tensor, + AscendC::GlobalTensor s_gm_tensor, + AscendC::GlobalTensor s_rope_gm_tensor, + AscendC::GlobalTensor mask_gm_tensor, + AscendC::LocalTensor dm32_ubuf_tensor, + AscendC::LocalTensor ll_ubuf_tensor, + AscendC::LocalTensor pm32_ubuf_tensor, + uint32_t n_idx, + uint32_t qk_n, + uint32_t qk_round_n, + uint32_t sub_m, + uint32_t mask_offset, + const uint32_t sub_n_loop, + const uint32_t cur_batch, + const uint32_t start_kv, + const uint32_t real_n_loop, + const uint32_t head_idx, + const uint32_t pm_flag_scalar, + uint32_t cur_q_seqlen, + uint32_t cur_kv_seqlen, + bool need_mask + ) + { + uint32_t sub_m_d128 = (sub_m + 127) / 128; // up aligned to 128 + uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; + float quantMax = (float)1 / (float)127; + WAIT_FLAG(V, MTE2, EVENT_ID2); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + DeQuantPerHeadImpl( + deq_scale_gm_tensor_q1[head_idx], + s_gm_tensor, + ls32_quant_ubuf_tensor, ls32_quant_ubuf_tensor.template ReinterpretCast(), + descale_q1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, qk_n, qk_round_n, 0, 1); + gm_to_ub( + ls32_ubuf_tensor.template ReinterpretCast(), + s_rope_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * qk_round_n / FLOAT_BLOCK_SIZE, + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + AscendC::Add(ls32_ubuf_tensor, ls32_ubuf_tensor, ls32_quant_ubuf_tensor, sub_m * qk_round_n); // float + PIPE_BARRIER(V); + } else { + gm_to_ub( + ls32_ubuf_tensor.template ReinterpretCast(), + s_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * qk_round_n / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + + // TODO add mask type condition + if (mask_type == 3) { + uint32_t aligned_mask_copy_len = RoundUp(qk_n); // 16 + uint32_t mask_dst_stride = (qk_round_n - aligned_mask_copy_len) / BLOCK_SIZE; // 0 + + AscendC::DataCopyPad( + mask_ubuf_tensor, + mask_gm_tensor, + AscendC::DataCopyExtParams( + cur_q_seqlen, + qk_n * 2, + maxKVSeqLen * 2 - qk_n * 2, + mask_dst_stride, + 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0) + ); + } else if (need_mask && mask_type == 4) { + AscendC::DataCopy( + mask_ubuf_tensor, + mask_gm_tensor, + AscendC::DataCopyParams( + cur_q_seqlen, // blockCount + qk_round_n * 2 / 32, // blockLen, 2 is sizeof(half) + MASK_COLUMNS * 2 / 32 - qk_round_n * 2 / 32, // srcStride + 0 // dstStride + ) + ); + } + + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + + if (mask_type == 3 || (need_mask && mask_type == 4)) { + AscendC::Cast( + mask32_ubuf_tensor, + mask_ubuf_tensor, + AscendC::RoundMode::CAST_NONE, + cur_q_seqlen * qk_round_n); + } + } + + for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { + muls_v(ls32_ubuf_tensor[vadd_idx * FLOAT_VECTOR_SIZE], + ls32_ubuf_tensor[vadd_idx * FLOAT_VECTOR_SIZE], + tor, + sub_m, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride + ); + } + if (qk_n % FLOAT_VECTOR_SIZE > 0) { + __set_mask(qk_n % FLOAT_VECTOR_SIZE); + muls_v(ls32_ubuf_tensor[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + ls32_ubuf_tensor[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tor, + sub_m, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + qk_round_n / FLOAT_BLOCK_SIZE, // dstRepeatStride + qk_round_n / FLOAT_BLOCK_SIZE // srcRepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + + if constexpr (tilingKeyType != TilingKeyType::TILING_INT8_DATA) { + if (mask_type == 3 || (need_mask && mask_type == 4)) { + uint32_t cur_compute_head_num = sub_m / cur_q_seqlen; + for (uint32_t i = 0; i < cur_compute_head_num; i++) { + Add( + ls32_ubuf_tensor[cur_q_seqlen * qk_round_n * i], + ls32_ubuf_tensor[cur_q_seqlen * qk_round_n * i], + mask32_ubuf_tensor, + cur_q_seqlen * qk_round_n + ); + } + PIPE_BARRIER(V); + } + } + + // *** lm = rowmax(ls) + ReduceMaxRepeatM(lm32_ubuf_tensor, ls32_ubuf_tensor, lp32_ubuf_tensor, sub_m, qk_n, qk_round_n); + // ReduceMaxChange(lm32_ubuf_tensor, ls32_ubuf_tensor, tv32_ubuf_tensor, round_sub_m, qk_n, qk_round_n); + if (n_idx != 0) { + // *** hm = vmax(lm, gm) + max_v(hm32_ubuf_tensor, + lm32_ubuf_tensor, + gm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + // *** dm = gm - hm + sub_v(dm32_ubuf_tensor, + gm32_ubuf_tensor, + hm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + } else { + // *** hm = lm + ub_to_ub( + hm32_ubuf_tensor, + lm32_ubuf_tensor, + 0, // sid + 1, // nBurst + round_sub_m / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + PIPE_BARRIER(V); + } + // *** gm = hm + ub_to_ub( + gm32_ubuf_tensor, + hm32_ubuf_tensor, + 0, // sid + 1, // nBurst + round_sub_m / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + PIPE_BARRIER(V); + // *** hm_block = expand_to_block(hm) + + // *** ls = ls - hm_block + TensorSubValueRepeatM(ls32_ubuf_tensor, ls32_ubuf_tensor, + hm32_ubuf_tensor, tv32_ubuf_tensor, + sub_m, round_sub_m, qk_n, qk_round_n); + // *** ls = exp(ls) + exp_v(ls32_ubuf_tensor, + ls32_ubuf_tensor, + (sub_m * qk_round_n + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + // *** lp = castfp32to16(ls) + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + sub_v(pm32_ubuf_tensor, + lm32_ubuf_tensor, + hm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + exp_v(pm32_ubuf_tensor, + pm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + muls_v(pm32_ubuf_tensor, + pm32_ubuf_tensor, + quantMax, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + brcb_v( + tv32_ubuf_tensor.ReinterpretCast(), + pm32_ubuf_tensor.ReinterpretCast(), + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + QuantPerTokenImpl(lp_ubuf_tensor, ls32_ubuf_tensor, tv32_ubuf_tensor, sub_m, qk_n, qk_round_n, 1); + } else { + conv_v(lp_ubuf_tensor, + ls32_ubuf_tensor, + (sub_m * qk_round_n + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + } + SET_FLAG(V, MTE3, EVENT_ID0); + WAIT_FLAG(V, MTE3, EVENT_ID0); + ub_to_gm( + p_gm_tensor, + lp_ubuf_tensor, + 0, // sid + 1, // nBurst + sub_m * qk_round_n * T_BLOCK_OFFSET / T_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + + // *** ll = rowsum(ls32) + ReduceSumRepeatM(ll_ubuf_tensor, ls32_ubuf_tensor, sub_m, qk_n, qk_round_n); + SET_FLAG(V, MTE2, EVENT_ID2); + PIPE_BARRIER(V); + } + + __aicore__ __attribute__((always_inline)) inline void SoftmaxStage2MLAHeadLoop( + AscendC::GlobalTensor o_tmp_gm_tensor, + AscendC::GlobalTensor go_gm_tensor, + AscendC::GlobalTensor o_gm_tensor, + AscendC::LocalTensor dm32_ubuf_tensor, + AscendC::LocalTensor ll_ubuf_tensor, + AscendC::LocalTensor pm32_ubuf_tensor, + uint32_t n_idx, + uint32_t n_loop, + uint32_t qk_n, + uint32_t qk_round_n, + uint32_t sub_m, + uint64_t o_offset, + uint32_t head_idx, + uint32_t pm_flag_scalar, + uint32_t head_loop, + uint32_t head_loop_idx, + uint32_t q_seq_len, + uint32_t sub_head_num, + uint32_t cur_head_num, + uint32_t numhead_per_process, + uint32_t head_res_row_num, + uint32_t head_start_sblock_idx, + uint32_t tail_res_row_num + ) + { + uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; + WAIT_FLAG(V, MTE2, EVENT_ID0); + if (n_idx != 0) { + gm_to_ub( + lo_ubuf_tensor.template ReinterpretCast(), + o_tmp_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + DeQuantPerHeadImpl( + deq_scale_gm_tensor_k1[head_idx], + o_tmp_gm_tensor, + lo_ubuf_tensor, lo_ubuf_tensor.template ReinterpretCast(),// lo_ubuf_tensor use the same ptr + descale_k1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, round_v, round_v, 1, 0); + } + } + SetVectorMask((uint64_t)-1, (uint64_t)-1); + WAIT_FLAG(MTE3, MTE2, EVENT_ID4); + if (n_idx != 0) { + // *** dm = exp(dm) + if (head_loop_idx == 0) { + exp_v(dm32_ubuf_tensor, + dm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + // *** gl = dm * gl + mul_v(gl32_ubuf_tensor, + dm32_ubuf_tensor, + gl32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + // *** gl = ll + gl + add_v(gl32_ubuf_tensor, + gl32_ubuf_tensor, + ll_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + } + SetVectorMask((uint64_t)-1, (uint64_t)-1); + brcb_v(tv32_ubuf_tensor.ReinterpretCast(), + dm32_ubuf_tensor.ReinterpretCast(), + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + if (head_loop > 1) { + gm_to_ub( + go32_ubuf_tensor, + go_gm_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + + // *** go = go * dm_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { + mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (__v % FLOAT_VECTOR_SIZE > 0) { + __set_mask(__v % FLOAT_VECTOR_SIZE); + mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + // *** go = lo + go + add_v(go32_ubuf_tensor, + go32_ubuf_tensor, + lo_ubuf_tensor, + (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + } else { + // *** gl = ll + if (head_loop_idx == 0) { + ub_to_ub( + gl32_ubuf_tensor, + ll_ubuf_tensor, + 0, // sid + 1, // nBurst + 64 / FLOAT_BLOCK_SIZE, // lenBurst + // round_sub_m / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + PIPE_BARRIER(V); + } + + gm_to_ub( + go32_ubuf_tensor.template ReinterpretCast(), + o_tmp_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { + DeQuantPerHeadImpl( + deq_scale_gm_tensor_k1[head_idx], + o_tmp_gm_tensor, + go32_ubuf_tensor, go32_ubuf_tensor.template ReinterpretCast(), + descale_k1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, round_v, round_v, 1, 0); + } else { + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + } + SET_FLAG(V, MTE2, EVENT_ID0); + + if (n_idx == n_loop - 1) { + // *** gl_block = expand_to_block(gl) + brcb_v(tv32_ubuf_tensor.ReinterpretCast(), + gl32_ubuf_tensor.ReinterpretCast()[head_loop_idx * 16], + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + // *** go = go / gl_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vdiv_idx = 0; vdiv_idx < __v / FLOAT_VECTOR_SIZE; ++vdiv_idx) { + div_v(go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (__v % FLOAT_VECTOR_SIZE > 0) { + __set_mask(__v % FLOAT_VECTOR_SIZE); + div_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); // fix hidden_size=96 + } + PIPE_BARRIER(V); + + // *** go = castfp32to16(go) + conv_v(go_ubuf_tensor, + go32_ubuf_tensor, + (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + SET_FLAG(V, MTE3, EVENT_ID0); + WAIT_FLAG(V, MTE3, EVENT_ID0); + + uint32_t inner_o_gm_offset = 0; + uint32_t inner_go_ubuf_offset = 0; + + if (head_res_row_num != 0) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset + q_heads * __v * head_start_sblock_idx], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams( + head_res_row_num, // blockCount + __v * 2, // blockLen + 0, // srcStride + __v * (q_heads - 1) * 2, // dstStride + 0 // rsv + ) + ); + inner_o_gm_offset += __v; + inner_go_ubuf_offset += head_res_row_num * __v; + } + + for (uint32_t i = 0; i < numhead_per_process; i++) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams( + q_seq_len, // blockCount + __v * 2, // blockLen + 0, // srcStride + __v * (q_heads - 1) * 2, // dstStride + 0 // rsv + ) + ); + inner_o_gm_offset += __v; + inner_go_ubuf_offset += q_seq_len * __v; + } + + if (tail_res_row_num != 0) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams( + tail_res_row_num, // blockCount + __v * 2, // blockLen + 0, // srcStride + __v * (q_heads - 1) * 2, // dstStride + 0 // rsv + ) + ); + } + // ********************* move O to GM ************************ + if constexpr (IS_RING) { + uint32_t lenBurst = sizeof(OUT_DTYPE); + ln_v(lse32_ubuf_tensor, + gl32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + add_v(lse32_ubuf_tensor, + lse32_ubuf_tensor, + gm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + conv_v(lse_conv_ubuf_tensor, + lse32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + SET_FLAG(V, MTE3, EVENT_ID1); + WAIT_FLAG(V, MTE3, EVENT_ID1); + // copyout lse + ub_to_gm_align( + lse_gm_tensor[(int64_t)(o_offset / __k)], + lse_conv_ubuf_tensor, + 0, // sid + 1, // nBurst + lenBurst * sub_m * head_loop, // lenBurst + 0, // leftPaddingNum + 0, // rightPaddingNum + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE3, V, EVENT_ID1); + WAIT_FLAG(MTE3, V, EVENT_ID1); + } + + } else if (head_loop > 1) { + SET_FLAG(V, MTE3, EVENT_ID5); + WAIT_FLAG(V, MTE3, EVENT_ID5); + ub_to_gm( + go_gm_tensor, + go32_ubuf_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + } + SET_FLAG(MTE3, MTE2, EVENT_ID4); + } + + __aicore__ __attribute__((always_inline)) inline void SoftmaxStage2MLAHeadLoopTP1( + AscendC::GlobalTensor o_tmp_gm_tensor, + AscendC::GlobalTensor go_gm_tensor, + AscendC::GlobalTensor o_gm_tensor, + AscendC::LocalTensor dm32_ubuf_tensor, + AscendC::LocalTensor ll_ubuf_tensor, + AscendC::LocalTensor pm32_ubuf_tensor, + uint32_t n_idx, + uint32_t n_loop, + uint32_t qk_n, + uint32_t qk_round_n, + uint32_t sub_m, + uint64_t o_offset, + uint32_t head_idx, + uint32_t pm_flag_scalar, + uint32_t head_loop, + uint32_t head_loop_idx, + uint32_t q_seq_len, + uint32_t sub_head_num, + uint32_t cur_head_num, + uint32_t numhead_per_process, + uint32_t head_res_row_num, + uint32_t head_start_sblock_idx, + uint32_t tail_res_row_num + ) + { + uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; + WAIT_FLAG(V, MTE2, EVENT_ID0); + if (n_idx != 4) { + gm_to_ub( + lo_ubuf_tensor.template ReinterpretCast(), + o_tmp_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + SetVectorMask((uint64_t)-1, (uint64_t)-1); + WAIT_FLAG(MTE3, MTE2, EVENT_ID4); + if (n_idx != 4) { + // expand_to_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + brcb_v(tv32_ubuf_tensor.ReinterpretCast(), + dm32_ubuf_tensor.ReinterpretCast(), + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + if (head_loop > 1) { + gm_to_ub( + go32_ubuf_tensor, + go_gm_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + + // *** go = go * dm_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { + mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (__v % FLOAT_VECTOR_SIZE > 0) { + __set_mask(__v % FLOAT_VECTOR_SIZE); + mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + PIPE_BARRIER(V); + // *** go = lo + go + add_v(go32_ubuf_tensor, + go32_ubuf_tensor, + lo_ubuf_tensor, + (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + } else { + // *** go = lo + + gm_to_ub( + go32_ubuf_tensor.template ReinterpretCast(), + o_tmp_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + SET_FLAG(V, MTE2, EVENT_ID0); + + if (n_idx + 4 > n_loop + 4 - 1) { + // *** gl_block = expand_to_block(gl), 存放于 tv + brcb_v(tv32_ubuf_tensor.ReinterpretCast(), + gl32_ubuf_tensor.ReinterpretCast()[head_loop_idx * 16], + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + // *** go = go / gl_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vdiv_idx = 0; vdiv_idx < __v / FLOAT_VECTOR_SIZE; ++vdiv_idx) { + div_v(go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (__v % FLOAT_VECTOR_SIZE > 0) { + __set_mask(__v % FLOAT_VECTOR_SIZE); + div_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); // fix hidden_size=96 + } + PIPE_BARRIER(V); + + // *** go = castfp32to16(go) + conv_v(go_ubuf_tensor, + go32_ubuf_tensor, + (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + SET_FLAG(V, MTE3, EVENT_ID0); + WAIT_FLAG(V, MTE3, EVENT_ID0); + + uint32_t inner_o_gm_offset = 0; + uint32_t inner_go_ubuf_offset = 0; + + if (head_res_row_num != 0) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset + q_heads * __v * head_start_sblock_idx], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams( + head_res_row_num, // blockCount + __v * 2, // blockLen + 0, // srcStride + __v * (q_heads - 1) * 2, // dstStride + 0 // rsv + ) + ); + inner_o_gm_offset += __v; + inner_go_ubuf_offset += head_res_row_num * __v; + } + + for (uint32_t i = 0; i < numhead_per_process; i++) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams( + q_seq_len, // blockCount + __v * 2, // blockLen + 0, // srcStride + __v * (q_heads - 1) * 2, // dstStride + 0 // rsv + ) + ); + inner_o_gm_offset += __v; + inner_go_ubuf_offset += q_seq_len * __v; + } + + if (tail_res_row_num != 0) { + AscendC::DataCopyPad( + o_gm_tensor[inner_o_gm_offset], + go_ubuf_tensor[inner_go_ubuf_offset], + AscendC::DataCopyExtParams( + tail_res_row_num, // blockCount + __v * 2, // blockLen + 0, // srcStride + __v * (q_heads - 1) * 2, // dstStride + 0 // rsv + ) + ); + } + // ********************* move O to GM ************************ + if constexpr (IS_RING) { + uint32_t lenBurst = sizeof(OUT_DTYPE); + ln_v(lse32_ubuf_tensor, + gl32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 8, // dstRepeatStride + 8 // srcRepeatStride + ); + PIPE_BARRIER(V); + add_v(lse32_ubuf_tensor, + lse32_ubuf_tensor, + gm32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + conv_v(lse_conv_ubuf_tensor, + lse32_ubuf_tensor, + sub_m_d64, // repeat + 1, // dstBlockStride + 1, // srcBlockStride + 4, // dstRepeatStride + 8 // srcRepeatStride + ); + SET_FLAG(V, MTE3, EVENT_ID1); + WAIT_FLAG(V, MTE3, EVENT_ID1); + // copyout lse + ub_to_gm_align( + lse_gm_tensor[(int64_t)(o_offset / __k)], + lse_conv_ubuf_tensor, + 0, // sid + 1, // nBurst + lenBurst * sub_m * head_loop, // lenBurst + 0, // leftPaddingNum + 0, // rightPaddingNum + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE3, V, EVENT_ID1); + WAIT_FLAG(MTE3, V, EVENT_ID1); + } + } + else if (head_loop > 1) { + SET_FLAG(V, MTE3, EVENT_ID5); + WAIT_FLAG(V, MTE3, EVENT_ID5); + ub_to_gm( + go_gm_tensor, + go32_ubuf_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + } + SET_FLAG(MTE3, MTE2, EVENT_ID4); + } - // other tasks - { - uint32_t cores_per_seq = 4; - // TODO: check whether need extra sync between aic and aiv for used s_gm - uint32_t process = totalTaskNumRound + block_num / 2 + block_idx / cores_per_seq; - if (process < totalTaskNum) { - uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; - uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - - uint32_t q_seqlen = 1; - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); - uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; - uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; - uint32_t kv_seqlen_process = (block_idx % cores_per_seq == cores_per_seq - 1) ? - (kv_seqlen - kv_seqlen_align * (cores_per_seq - 1)) : kv_seqlen_align; - - if (kv_seqlen > 0 && kv_seqlen_process > 0) { - uint32_t start_head = 0; - uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen_process; - uint32_t cur_head_num = q_heads; - uint32_t cur_nIndx = 0; // no use, follow the previous code - // customized InnerRunVectorChange for tail processing - TailInnerRunVectorChangeTP1(start_head, cur_q_seq_len, cur_kv_seqlen, cur_head_num, offset_tiling, 512, embed_split_loop_v_former); - // Sync all vector cores - AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); - AscendC::CrossCoreWaitFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); - - // 8 vector cores - if (block_idx % cores_per_seq < 4) { - TailInnerGatherVectorTP1(start_head + cur_head_num / 4 * (block_idx % cores_per_seq), cur_q_seq_len, cur_head_num / 4, block_idx - (block_idx % cores_per_seq), cores_per_seq, offset_tiling); - } - - } - else{ - AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); - AscendC::CrossCoreWaitFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); - } - } - else{ - AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); - AscendC::CrossCoreWaitFlag<0x0, PIPE_MTE3>(TAIL_OPTIMIZATION_SYNC); - } - } + __aicore__ __attribute__((always_inline)) inline void TailSoftmaxStage2MLAHeadLoopTP1( + AscendC::GlobalTensor o_tmp_gm_tensor, + AscendC::GlobalTensor go_gm_tensor, + AscendC::GlobalTensor gl_gm_tensor, + AscendC::GlobalTensor gm_gm_tensor, + AscendC::LocalTensor dm32_ubuf_tensor, + AscendC::LocalTensor go32_ubuf_tensor, + AscendC::LocalTensor gl32_ubuf_tensor, + AscendC::LocalTensor gm32_ubuf_tensor, + uint32_t n_idx, + uint32_t n_loop, + uint32_t qk_n, + uint32_t qk_round_n, + uint32_t sub_m, + uint64_t o_offset, + uint32_t head_idx, + uint32_t head_loop, + uint32_t head_loop_idx, + uint32_t q_seq_len, + uint32_t sub_head_num, + uint32_t cur_head_num, + uint32_t numhead_per_process + ) + { + uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; + WAIT_FLAG(V, MTE2, EVENT_ID0); + if (n_idx != 4) { + gm_to_ub( + lo_ubuf_tensor.template ReinterpretCast(), + o_tmp_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + SetVectorMask((uint64_t)-1, (uint64_t)-1); + WAIT_FLAG(MTE3, MTE2, EVENT_ID4); + if (n_idx != 4) { + // expand_to_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + brcb_v(tv32_ubuf_tensor.ReinterpretCast(), + dm32_ubuf_tensor.ReinterpretCast(), + 1, // dstBlockStride + 8, // dstRepeatStride + round_sub_m / FLOAT_BLOCK_SIZE // repeat + ); + PIPE_BARRIER(V); + if (head_loop > 1) { + gm_to_ub( + go32_ubuf_tensor, + go_gm_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); + } + + // *** go = go * dm_block + SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { + mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + } + if (__v % FLOAT_VECTOR_SIZE > 0) { + __set_mask(__v % FLOAT_VECTOR_SIZE); + mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tv32_ubuf_tensor, + sub_m, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 0, // src1BlockStride + round_v / FLOAT_BLOCK_SIZE, // dstRepeatStride + round_v / FLOAT_BLOCK_SIZE, // src0RepeatStride + 1 // src1RepeatStride + ); + SetVectorMask((uint64_t)-1, (uint64_t)-1); } + PIPE_BARRIER(V); + // *** go = lo + go + add_v(go32_ubuf_tensor, + go32_ubuf_tensor, + lo_ubuf_tensor, + (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat + 1, // dstBlockStride + 1, // src0BlockStride + 1, // src1BlockStride + 8, // dstRepeatStride + 8, // src0RepeatStride + 8 // src1RepeatStride + ); + PIPE_BARRIER(V); + } else { + // *** go = lo + + gm_to_ub( + go32_ubuf_tensor.template ReinterpretCast(), + o_tmp_gm_tensor, + 0, // sid + 1, // nBurst + sub_m * round_v / FLOAT_BLOCK_SIZE, // lenBurst + 0, // srcGap + 0 // dstGap + ); + SET_FLAG(MTE2, V, EVENT_ID0); + WAIT_FLAG(MTE2, V, EVENT_ID0); } + SET_FLAG(V, MTE2, EVENT_ID0); - PlatformWaitVectorPipeSyncTP1(); + if (n_idx + 4 > n_loop + 4 - 1) { + // The last step to process the o with dividing and copyout + // TODO: Maybe the following two don't need waiting + // Copyout gl32_ubuf_tensor to gl_gm_tensor + ub_to_gm( + gl_gm_tensor, + gl32_ubuf_tensor, + 0, + 1, + sub_m / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + + // Copyout rowmax to global gm32_ubuf_tensor + ub_to_gm( + gm_gm_tensor, + gm32_ubuf_tensor, + 0, + 1, + sub_m / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + // Copyout go32_ubuf_tensor to go_gm_tensor + // This is needed to wait for former calculation + SET_FLAG(V, MTE3, EVENT_ID5); + WAIT_FLAG(V, MTE3, EVENT_ID5); + ub_to_gm( + go_gm_tensor, + go32_ubuf_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + + } + else if (head_loop > 1) { + SET_FLAG(V, MTE3, EVENT_ID5); + WAIT_FLAG(V, MTE3, EVENT_ID5); + ub_to_gm( + go_gm_tensor, + go32_ubuf_tensor, + 0, + 1, + sub_m * round_v / FLOAT_BLOCK_SIZE, + 0, + 0 + ); + } + SET_FLAG(MTE3, MTE2, EVENT_ID4); + PIPE_BARRIER(ALL); } -private: - // ====== AIV refactor: context struct + sub-functions ====== - // VectorContext / InitVectorContext / VectorTP1Context / InitVectorTP1Context - // 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） + __aicore__ __attribute__((always_inline)) inline void InnerRunVectorChange( + uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, + uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, + uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) + { + uint32_t addr_o_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 4 + offset_tiling)); + uint32_t addr_o_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 5 + offset_tiling)); + uint64_t addr_o_scalar = (uint64_t)(((uint64_t)addr_o_high32) << 32 | addr_o_loww32); + + uint32_t addr_mask_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 6 + offset_tiling)); + uint32_t addr_mask_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 7 + offset_tiling)); + uint64_t addr_mask_scalar = (uint64_t)(((uint64_t)addr_mask_high32) << 32 | addr_mask_loww32); + + uint32_t mask_offset = addr_mask_scalar; + + uint32_t pp_n_scalar = block_size; // 64 + uint32_t sub_n_loop = pp_n_scalar / block_size; + uint32_t real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; - // [AIV 第三层平台函数 + 第二层业务函数] 已迁移至 aiv_arch32.h / aiv_bs.h - #include "multi_latent_attention_aiv_arch32.h" - #include "multi_latent_attention_aiv_bs.h" + uint32_t n_loop = (cur_kv_seqlen + pp_n_scalar - 1) / pp_n_scalar; - // ReduceMaxRepeatM - // 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） + uint32_t qk_n = pp_n_scalar; + uint32_t qk_round_n = RoundUp(qk_n); + uint32_t qk_n_2 = pp_n_scalar; + uint32_t qk_round_n_2 = RoundUp(qk_n_2); - // ReduceSumRepeatM - // 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） + // split head num to two vectors + uint32_t sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; // 16 + uint32_t sub_m = sub_head_num * cur_q_seqlen; // 16 * 3 = 48 + + uint32_t head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; // not used + + o_offset = addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; // for NSD -> SND + + uint32_t sub_m_d128 = (sub_m + 127) / 128; // up aligned to 128 + uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; + + uint32_t start_kv = 0; + /* if tail length smalller than q_len - 1, then need to mask the last two tile*/ + uint32_t tail_len = cur_kv_seqlen - (n_loop - 1) * pp_n_scalar; + bool prev_tail_mask = (n_loop > 1 && tail_len < cur_q_seqlen - 1); + for (uint32_t n_idx = 0; n_idx < n_loop + 1; n_idx++) { + if (n_idx != n_loop) { + bool need_mask = false; + uint32_t mask_start_offset = 0; + if (n_idx == (n_loop - 2)) { + need_mask = prev_tail_mask; + mask_start_offset = need_mask ? (tail_len + MASK_COLUMNS - 1) * MASK_COLUMNS : 0; + } + if (n_idx == (n_loop - 1)) { + qk_n = (cur_kv_seqlen - n_idx * pp_n_scalar); + qk_round_n = RoundUp<16>(qk_n); + need_mask = true; + mask_start_offset = (qk_n - 1) * MASK_COLUMNS; + } + WaitFlagDev(QK_READY_DECODER); + /* ************ softmax1 stage1 ************* */ + WAIT_FLAG(MTE3, MTE2, EVENT_ID3); + if (sub_m > 0) { + if (mask_type == 3) { + mask_start_offset = mask_offset + n_idx * pp_n_scalar; + } + // input QK shape (sub_m, qk_round_n) + if (n_idx % 2 == 0){ + SoftmaxStage1( + p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + + (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n * T_BLOCK_OFFSET + (uint64_t)(n_idx % 2) * TMP_SIZE * T_BLOCK_OFFSET / 2], + s_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + + (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + s_rope_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + + (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], + mask_gm_tensor[mask_start_offset], + dm32_ubuf_tensor, ll_ubuf_tensor, pm32_ubuf_tensor, + n_idx, qk_n, qk_round_n, sub_m, 0, sub_n_loop, cur_batch, start_kv, real_n_loop, head_idx, pm_flag_scalar1, cur_q_seqlen, cur_kv_seqlen, need_mask + ); + } else { + SoftmaxStage1( + p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + + (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n * T_BLOCK_OFFSET + + TMP_SIZE * T_BLOCK_OFFSET / 2], + s_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + + (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + + TMP_SIZE_DECODER / 2], + s_rope_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + + (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + + TMP_SIZE_DECODER / 2], + mask_gm_tensor[mask_start_offset], + dm32_stage2_ubuf_tensor, ll_stage2_ubuf_tensor, pm32_ubuf_stage2_tensor, + n_idx, qk_n, qk_round_n, sub_m, 0, sub_n_loop, cur_batch, start_kv, real_n_loop, head_idx, pm_flag_scalar2, cur_q_seqlen, cur_kv_seqlen, need_mask + ); + } + } + FftsCrossCoreSync(SOFTMAX_READY_DECODER); - // TensorSubValueRepeatM / TensorDivRepeatM / TensorMulRepeatM / DeQuantPerHeadImpl / QuantPerTokenImpl - // 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） + SET_FLAG(MTE3, MTE2, EVENT_ID3); + } + /* ************ softmax2 stage1 ************* */ - // SoftmaxStage2MLAHeadLoop / SoftmaxStage2MLAHeadLoopTP1 / TailSoftmaxStage2MLAHeadLoopTP1 - // 已迁移至 multi_latent_attention_aiv_bs.h（业务逻辑层） + uint32_t process_row_num = 16; + uint32_t numhead_per_process = process_row_num / cur_q_seqlen; + if (n_idx != 0) { + if (n_idx == n_loop) { + qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); + qk_round_n_2 = RoundUp(qk_n_2); + } + WaitFlagDev(UPDATE_READY_DECODER); + if (sub_m > 0) { + uint32_t head_loop = (sub_m + process_row_num - 1) / process_row_num; + + uint32_t head_res_row_num = 0; + uint32_t head_start_sblock_idx = 0; + uint32_t tail_res_row_num = 0; + + for (uint32_t head_loop_idx = 0; head_loop_idx < head_loop; ++head_loop_idx) { + uint32_t head_offset = head_loop_idx * process_row_num * round_v; + uint32_t cur_sub_m = head_loop_idx == (head_loop - 1) ? sub_m - head_loop_idx * process_row_num : process_row_num; // 15 or 3 + + // complete head num + head_start_sblock_idx = tail_res_row_num; + head_res_row_num = (cur_q_seqlen - tail_res_row_num) % cur_q_seqlen; + uint32_t cur_numhead_per_process = (cur_sub_m - head_res_row_num) / cur_q_seqlen; + tail_res_row_num = cur_sub_m - cur_numhead_per_process * cur_q_seqlen - head_res_row_num; + + uint32_t out_o_offset = head_loop_idx * numhead_per_process * round_v; // modified, round_v = 512 + + SoftmaxStage2MLAHeadLoop( + o_tmp_gm_tensor[(uint64_t)(block_idx * TMP_SIZE * 2 + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset + ((n_idx - 1) % 2) * TMP_SIZE)], + go_gm_tensor[(uint64_t)(block_idx * TMP_SIZE + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset)], + o_gm_tensor[(uint64_t)(o_offset + out_o_offset)], + dm32_ubuf_tensor[(uint64_t)((n_idx - 1) % 2 * 128 + head_loop_idx * process_row_num)], + ll_ubuf_tensor[(uint64_t)((n_idx - 1) % 2 * 256 + head_loop_idx * process_row_num)], + pm32_ubuf_tensor[(uint64_t)((n_idx - 1) % 2 * 128 + head_loop_idx * process_row_num)], + n_idx - 1, n_loop, qk_n_2, RoundUp(qk_round_n_2), cur_sub_m, o_offset, + head_idx + head_loop_idx * process_row_num, + pm_flag_scalar1, head_loop, head_loop_idx, cur_q_seqlen, sub_head_num, cur_head_num, + cur_numhead_per_process, + head_res_row_num, head_start_sblock_idx, tail_res_row_num); + } + } + } + } + } __aicore__ __attribute((always_inline)) inline void SoftmaxGatherTP1( AscendC::GlobalTensor o_gm_tensor, @@ -1679,6 +4126,273 @@ class MLADecoderAiv{ AscendC::PipeBarrier(); } + __aicore__ __attribute__((always_inline)) inline void InnerRunVectorChangeTP1( + uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, + uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, + uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) + { + uint32_t prev_task = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); + uint64_t addr_o_scalar = prev_task * q_heads * embedding_size; + uint64_t addr_mask_scalar = 0; + uint32_t mask_offset = addr_mask_scalar; + + uint32_t pp_n_scalar = block_size; // 64 + uint32_t sub_n_loop = pp_n_scalar / block_size; + uint32_t real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; + + uint32_t n_loop = (cur_kv_seqlen + pp_n_scalar - 1) / pp_n_scalar; + + uint32_t qk_n = pp_n_scalar; + uint32_t qk_round_n = RoundUp(qk_n); + + uint32_t qk_n_2 = pp_n_scalar; + uint32_t qk_round_n_2 = RoundUp(qk_n_2); + + // split head num to two vectors + uint32_t sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; // 16 + uint32_t sub_m = sub_head_num * cur_q_seqlen; // 16 * 3 = 48 + + uint32_t head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; // not used + + o_offset = addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; // for NSD -> SND + + uint32_t sub_m_d128 = (sub_m + 127) / 128; // up aligned to 128 + uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; + + uint32_t start_kv = 0; + uint32_t s_block_stack = 4; + uint32_t m_slice = FLOAT_VECTOR_SIZE / s_block_stack; + uint32_t m_end = (sub_m + m_slice - 1) / m_slice; + for (uint32_t n_idx = 0; n_idx < n_loop + s_block_stack; n_idx += s_block_stack) { + if (n_idx < n_loop) { + if (n_idx + s_block_stack > n_loop - 1) { + qk_n = (cur_kv_seqlen - n_idx * pp_n_scalar); + } else { + qk_n = pp_n_scalar * s_block_stack; + } + qk_round_n = RoundUp<16>(qk_n); + if (sub_m == 0) { + WaitFlagDev(QK_READY_DECODER); + } + uint32_t pingpong_flag = 0; + for (uint32_t m_ind = 0; m_ind < m_end; m_ind++) { + uint32_t row_offset = m_ind * m_slice; + uint32_t curr_m = m_ind == m_end - 1 ? sub_m - row_offset : m_slice; + uint32_t s_ub_offset = pingpong_flag * 8192; + uint32_t p_gm_offset = (uint64_t)block_idx * TMP_SIZE * 2 + + (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + row_offset * qk_round_n + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE; + uint32_t s_gm_offset = (int64_t)block_idx * TMP_SIZE_DECODER * 4 + + (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + row_offset * qk_round_n + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE_DECODER * 2; + if (m_ind == 0) { + WaitFlagDev(QK_READY_DECODER); + } + if (curr_m == 0) { + continue; + } + OnlineSoftmaxStage1 ( + ls32_ubuf_tensor[s_ub_offset], + mask_ubuf_tensor, + mask_ubuf_tensor.template ReinterpretCast(), + lm32_ubuf_tensor[row_offset], + hm32_ubuf_tensor[row_offset], + gm32_ubuf_tensor[row_offset], + dm32_ubuf_tensor[((n_idx / s_block_stack) % 2) * UB_FLOAT_LINE_SIZE + row_offset], + ls32_ubuf_tensor[s_ub_offset], + ll_ubuf_tensor[row_offset], + gl32_ubuf_tensor[row_offset], + lp_ubuf_tensor[s_ub_offset * 2], + tv32_ubuf_tensor, + s_gm_tensor[s_gm_offset], + p_gm_tensor[p_gm_offset], + n_idx == 0, this->tor, + curr_m, qk_n, qk_round_n, pingpong_flag + ); + pingpong_flag = 1 - pingpong_flag; + } + FftsCrossCoreSync(SOFTMAX_READY_DECODER); + } + /* ************ softmax2 stage1 ************* */ + // PIPE_BARRIER(ALL); + uint32_t process_row_num = 16; + uint32_t numhead_per_process = process_row_num / cur_q_seqlen; + + if (n_idx >= s_block_stack) { + if (n_idx == n_loop) { + qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); + qk_round_n_2 = RoundUp(qk_n_2); + } + WaitFlagDev(UPDATE_READY_DECODER); + if (sub_m > 0) { + uint32_t head_loop = (sub_m + process_row_num - 1) / process_row_num; + + uint32_t head_res_row_num = 0; + uint32_t head_start_sblock_idx = 0; + uint32_t tail_res_row_num = 0; + + for (uint32_t head_loop_idx = 0; head_loop_idx < head_loop; ++head_loop_idx) { + uint32_t head_offset = head_loop_idx * process_row_num * round_v; + uint32_t cur_sub_m = head_loop_idx == (head_loop - 1) ? sub_m - head_loop_idx * process_row_num : process_row_num; // 15 or 3 + + // complete head num + head_start_sblock_idx = tail_res_row_num; + head_res_row_num = (cur_q_seqlen - tail_res_row_num) % cur_q_seqlen; + uint32_t cur_numhead_per_process = (cur_sub_m - head_res_row_num) / cur_q_seqlen; + tail_res_row_num = cur_sub_m - cur_numhead_per_process * cur_q_seqlen - head_res_row_num; + + uint32_t out_o_offset = head_loop_idx * numhead_per_process * round_v; // modified, round_v = 512 + + SoftmaxStage2MLAHeadLoopTP1( + o_tmp_gm_tensor[(uint64_t)(block_idx * TMP_SIZE * 2 + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset + ((n_idx / s_block_stack - 1) % 2) * TMP_SIZE)], + go_gm_tensor[(uint64_t)(block_idx * TMP_SIZE + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset)], + o_gm_tensor[(uint64_t)(o_offset + out_o_offset)], + dm32_ubuf_tensor[(uint64_t)((n_idx / s_block_stack - 1) % 2 * UB_FLOAT_LINE_SIZE + head_loop_idx * process_row_num)], + ll_ubuf_tensor[(uint64_t)((n_idx / s_block_stack - 1) % 2 * 256 + head_loop_idx * process_row_num)], + pm32_ubuf_tensor, + n_idx, n_loop, qk_n_2, RoundUp(qk_round_n_2), cur_sub_m, o_offset, head_idx, + pm_flag_scalar1, head_loop, head_loop_idx, cur_q_seqlen, sub_head_num, cur_head_num, cur_numhead_per_process, + head_res_row_num, head_start_sblock_idx, tail_res_row_num); + } + } + } + } + } + + __aicore__ __attribute__((always_inline)) inline void TailInnerRunVectorChangeTP1( + uint32_t start_head, + uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, + uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) + { + uint32_t prev_task = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); + uint64_t addr_o_scalar = prev_task * q_heads * embedding_size; + uint64_t addr_mask_scalar = 0; + uint32_t mask_offset = addr_mask_scalar; + + uint32_t pp_n_scalar = block_size; // 64 + uint32_t sub_n_loop = pp_n_scalar / block_size; + uint32_t real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; + + uint32_t n_loop = (cur_kv_seqlen + pp_n_scalar - 1) / pp_n_scalar; + + uint32_t qk_n = pp_n_scalar; + uint32_t qk_round_n = RoundUp(qk_n); + + uint32_t qk_n_2 = pp_n_scalar; + uint32_t qk_round_n_2 = RoundUp(qk_n_2); + + // split head num to two vectors + uint32_t sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; // 16 + uint32_t sub_m = sub_head_num * cur_q_seqlen; // 16 * 3 = 48 + + uint32_t head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; // not used + + o_offset = addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; // for NSD -> SND + + uint32_t sub_m_d128 = (sub_m + 127) / 128; // up aligned to 128 + uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 + uint32_t round_sub_m = (sub_m + 15) / 16 * 16; + + uint32_t start_kv = 0; + uint32_t s_block_stack = 4; + uint32_t m_slice = FLOAT_VECTOR_SIZE / s_block_stack; + uint32_t m_end = (sub_m + m_slice - 1) / m_slice; + for (uint32_t n_idx = 0; n_idx < n_loop + s_block_stack; n_idx += s_block_stack) { + if (n_idx < n_loop) { + if (n_idx + s_block_stack > n_loop - 1) { + qk_n = (cur_kv_seqlen - n_idx * pp_n_scalar); + } else { + qk_n = pp_n_scalar * s_block_stack; + } + qk_round_n = RoundUp<16>(qk_n); + if (sub_m == 0) { + WaitFlagDev(QK_READY_DECODER); + } + uint32_t pingpong_flag = 0; + for (uint32_t m_ind = 0; m_ind < m_end; m_ind++) { + uint32_t row_offset = m_ind * m_slice; + uint32_t curr_m = m_ind == m_end - 1 ? sub_m - row_offset : m_slice; + uint32_t s_ub_offset = pingpong_flag * 8192; + uint32_t p_gm_offset = (uint64_t)block_idx * TMP_SIZE * 2 + + (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + row_offset * qk_round_n + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE; + uint32_t s_gm_offset = (int64_t)block_idx * TMP_SIZE_DECODER * 4 + + (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + row_offset * qk_round_n + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE_DECODER * 2; + if (m_ind == 0) { + WaitFlagDev(QK_READY_DECODER); + } + if (curr_m == 0) { + continue; + } + OnlineSoftmaxStage1 ( + ls32_ubuf_tensor[s_ub_offset], + mask_ubuf_tensor, + mask_ubuf_tensor.template ReinterpretCast(), + lm32_ubuf_tensor[row_offset], + hm32_ubuf_tensor[row_offset], + gm32_ubuf_tensor[row_offset], + dm32_ubuf_tensor[((n_idx / s_block_stack) % 2) * UB_FLOAT_LINE_SIZE + row_offset], + ls32_ubuf_tensor[s_ub_offset], + ll_ubuf_tensor[row_offset], + gl32_ubuf_tensor[row_offset], + lp_ubuf_tensor[s_ub_offset * 2], + tv32_ubuf_tensor, + s_gm_tensor[s_gm_offset], + p_gm_tensor[p_gm_offset], + n_idx == 0, this->tor, + curr_m, qk_n, qk_round_n, pingpong_flag + ); + pingpong_flag = 1 - pingpong_flag; + } + FftsCrossCoreSync(SOFTMAX_READY_DECODER); + } + /* ************ softmax2 stage1 ************* */ + // PIPE_BARRIER(ALL); + uint32_t process_row_num = 16; + uint32_t numhead_per_process = process_row_num / cur_q_seqlen; + + if (n_idx >= s_block_stack) { + if (n_idx == n_loop) { + qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); + qk_round_n_2 = RoundUp(qk_n_2); + } + WaitFlagDev(UPDATE_READY_DECODER); + if (sub_m > 0) { + uint32_t head_loop = (sub_m + process_row_num - 1) / process_row_num; + + uint32_t head_res_row_num = 0; + uint32_t head_start_sblock_idx = 0; + uint32_t tail_res_row_num = 0; + + for (uint32_t head_loop_idx = 0; head_loop_idx < head_loop; ++head_loop_idx) { + uint32_t head_offset = head_loop_idx * process_row_num * round_v; + uint32_t cur_sub_m = head_loop_idx == (head_loop - 1) ? sub_m - head_loop_idx * process_row_num : process_row_num; + + // complete head num + head_start_sblock_idx = tail_res_row_num; + head_res_row_num = (cur_q_seqlen - tail_res_row_num) % cur_q_seqlen; + uint32_t cur_numhead_per_process = (cur_sub_m - head_res_row_num) / cur_q_seqlen; + tail_res_row_num = cur_sub_m - cur_numhead_per_process * cur_q_seqlen - head_res_row_num; + + uint32_t out_o_offset = head_loop_idx * numhead_per_process * round_v; // round_v = 512 + + TailSoftmaxStage2MLAHeadLoopTP1( + o_tmp_gm_tensor[(uint64_t)(block_idx * TMP_SIZE * 2 + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset + ((n_idx / s_block_stack - 1) % 2) * TMP_SIZE)], + go_gm_tensor[(uint64_t)(block_idx * TMP_SIZE + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset)], + tmp_gm_tensor[(uint64_t)(block_idx * q_heads + start_head + sub_block_idx * cur_head_num * cur_q_seqlen / 2 + head_loop_idx * process_row_num)], + tmp_gm_tensor[(uint64_t)(block_num * q_heads + block_idx * q_heads + start_head + sub_block_idx * cur_head_num * cur_q_seqlen / 2 + head_loop_idx * process_row_num)], + dm32_ubuf_tensor[(uint64_t)((n_idx / s_block_stack - 1) % 2 * UB_FLOAT_LINE_SIZE + head_loop_idx * process_row_num)], + go32_ubuf_tensor, // no need for offset + gl32_ubuf_tensor[(uint64_t)(head_loop_idx * process_row_num)], + gm32_ubuf_tensor[(uint64_t)(head_loop_idx * process_row_num)], + n_idx, n_loop, qk_n_2, RoundUp(qk_round_n_2), cur_sub_m, o_offset, head_idx, + head_loop, head_loop_idx, cur_q_seqlen, sub_head_num, cur_head_num, cur_numhead_per_process + ); + } + } + } + } + + } + __aicore__ __attribute__((always_inline)) inline void TailInnerGatherVectorTP1( uint32_t start_head, uint32_t cur_q_seqlen, uint32_t cur_head_num, uint32_t start_block_idx, uint32_t cores_process, uint32_t offset_tiling) diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h deleted file mode 100644 index 8477643..0000000 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_arch32.h +++ /dev/null @@ -1,1089 +0,0 @@ -// Copyright 2025 The xLLM Authors. 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: -// -// https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE -// -// ==================== AIV 平台处理层（第三层）==================== -// 本文件仅被 MLADecoderAiv 类 include，包含 Vector 侧平台同步函数。 -// AIC 侧平台函数在 multi_latent_attention_arch32.h 中。 - -#pragma once - -// 平台函数：非TP1 SoftmaxStage1 前置管道同步 -// 在 SoftmaxStage1 调用之前执行：等待 QK 就绪 + 等待 MTE3→MTE2 管道同步 -__aicore__ __attribute__((always_inline)) inline void PlatformSoftmaxStage1PreSync() -{ - WaitFlagDev(QK_READY_DECODER); - WAIT_FLAG(MTE3, MTE2, EVENT_ID3); -} - -// 平台函数：非TP1 SoftmaxStage1 后置管道同步 -// 在 SoftmaxStage1 调用之后执行：核间同步 + 设置 MTE3→MTE2 管道 flag -__aicore__ __attribute__((always_inline)) inline void PlatformSoftmaxStage1PostSync() -{ - FftsCrossCoreSync(SOFTMAX_READY_DECODER); - SET_FLAG(MTE3, MTE2, EVENT_ID3); -} - -// 平台函数：Vector 非TP1 管道同步初始化（Run 方法头部 9 条 SET_FLAG） -__aicore__ __attribute__((always_inline)) inline void PlatformInitVectorPipeSync() -{ - SET_FLAG(MTE3, V, EVENT_ID0); - SET_FLAG(MTE3, MTE2, EVENT_ID0); - SET_FLAG(MTE3, MTE2, EVENT_ID2); - SET_FLAG(MTE3, MTE2, EVENT_ID3); - SET_FLAG(MTE3, MTE2, EVENT_ID4); - SET_FLAG(V, MTE2, EVENT_ID4); - SET_FLAG(V, MTE2, EVENT_ID0); - SET_FLAG(MTE3, V, EVENT_ID2); - SET_FLAG(V, MTE2, EVENT_ID2); -} - -// 平台函数：等待 Vector 非TP1 管道同步完成（Run 方法尾部 9 条 WAIT_FLAG） -__aicore__ __attribute__((always_inline)) inline void PlatformWaitVectorPipeSync() -{ - WAIT_FLAG(MTE3, V, EVENT_ID0); - WAIT_FLAG(MTE3, MTE2, EVENT_ID0); - WAIT_FLAG(MTE3, MTE2, EVENT_ID2); - WAIT_FLAG(MTE3, MTE2, EVENT_ID3); - WAIT_FLAG(MTE3, MTE2, EVENT_ID4); - WAIT_FLAG(V, MTE2, EVENT_ID0); - WAIT_FLAG(V, MTE2, EVENT_ID4); - WAIT_FLAG(MTE3, V, EVENT_ID2); - WAIT_FLAG(V, MTE2, EVENT_ID2); -} - -// 平台函数：MTE2→V EVENT_ID0 同步（SET+WAIT） -// 用于 DeQuantPerHeadImpl：src 搬入后通知 V 管道可读 -__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSyncEvent0() -{ - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); -} - -// 平台函数：MTE2→V EVENT_ID2 同步（SET+WAIT） -// 用于 DeQuantPerHeadImpl：deScale 搬入后通知 V 管道可做 online 乘法 -__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSyncEvent2() -{ - SET_FLAG(MTE2, V, EVENT_ID2); - WAIT_FLAG(MTE2, V, EVENT_ID2); -} - -// 平台函数：V→MTE2 管道等待（等待上一轮 V→MTE2 通知） -// 用于 SoftmaxStage1 LoadQKData 段开头：等待上一轮 QuantizeAndOutput 发出的 MTE2 可搬入通知 -__aicore__ __attribute__((always_inline)) inline void PlatformVToMte2Wait() -{ - WAIT_FLAG(V, MTE2, EVENT_ID2); -} - -// 平台函数：V 管道同步屏障 -// 用于 SoftmaxStage1 LoadQKData 段 INT8 路径：Add 之后确保数据可见 -__aicore__ __attribute__((always_inline)) inline void PlatformVPipeBarrier() -{ - PIPE_BARRIER(V); -} - -// 平台函数：Vector TP1 管道同步初始化（RunTP1 方法头部 10 条 SET_FLAG） -// 与非TP1 相比多出 SET_FLAG(MTE3, MTE2, EVENT_ID1) -__aicore__ __attribute__((always_inline)) inline void PlatformInitVectorPipeSyncTP1() -{ - SET_FLAG(MTE3, V, EVENT_ID0); - SET_FLAG(MTE3, MTE2, EVENT_ID0); - SET_FLAG(MTE3, MTE2, EVENT_ID1); - SET_FLAG(MTE3, MTE2, EVENT_ID2); - SET_FLAG(MTE3, MTE2, EVENT_ID3); - SET_FLAG(MTE3, MTE2, EVENT_ID4); - SET_FLAG(V, MTE2, EVENT_ID4); - SET_FLAG(V, MTE2, EVENT_ID0); - SET_FLAG(MTE3, V, EVENT_ID2); - SET_FLAG(V, MTE2, EVENT_ID2); -} - -// 平台函数：等待 Vector TP1 管道同步完成（RunTP1 方法尾部 10 条 WAIT_FLAG） -// 与非TP1 相比多出 WAIT_FLAG(MTE3, MTE2, EVENT_ID1) -__aicore__ __attribute__((always_inline)) inline void PlatformWaitVectorPipeSyncTP1() -{ - WAIT_FLAG(MTE3, V, EVENT_ID0); - WAIT_FLAG(MTE3, MTE2, EVENT_ID0); - WAIT_FLAG(MTE3, MTE2, EVENT_ID1); - WAIT_FLAG(MTE3, MTE2, EVENT_ID2); - WAIT_FLAG(MTE3, MTE2, EVENT_ID3); - WAIT_FLAG(MTE3, MTE2, EVENT_ID4); - WAIT_FLAG(V, MTE2, EVENT_ID0); - WAIT_FLAG(V, MTE2, EVENT_ID4); - WAIT_FLAG(MTE3, V, EVENT_ID2); - WAIT_FLAG(V, MTE2, EVENT_ID2); -} - -// 平台函数：向量广播乘法 go = go * dm_block -// 封装 brcb dm → mul 循环 → mask 尾部 → PIPE_BARRIER -// 用于 Stage2 段落B：go = go * dm_block（三函数完全一致） -__aicore__ __attribute__((always_inline)) inline void PlatformMulVectorByBroadcast( - AscendC::LocalTensor go32_ubuf_tensor, - AscendC::LocalTensor tv32_ubuf_tensor, - AscendC::LocalTensor dm32_ubuf_tensor, - uint32_t sub_m, - uint32_t round_sub_m, - uint32_t round_v, - uint32_t __v) -{ - // expand_to_block: dm → tv - SetVectorMask((uint64_t)-1, (uint64_t)-1); - brcb_v(tv32_ubuf_tensor, - dm32_ubuf_tensor.ReinterpretCast(), - 1, // dstBlockStride - 8, // dstRepeatStride - (uint8_t)(round_sub_m / FLOAT_BLOCK_SIZE) // repeat - ); - PIPE_BARRIER(V); - - // go = go * dm_block (full vector iterations) - SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vmul_idx = 0; vmul_idx < __v / FLOAT_VECTOR_SIZE; ++vmul_idx) { - mul_v(go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[vmul_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor.template ReinterpretCast(), - (uint8_t)sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // dstRepeatStride - (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // src0RepeatStride - 1 // src1RepeatStride - ); - } - // tail mask iteration - if (__v % FLOAT_VECTOR_SIZE > 0) { - __set_mask(__v % FLOAT_VECTOR_SIZE); - mul_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor.template ReinterpretCast(), - (uint8_t)sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // dstRepeatStride - (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); -} - -// 平台函数：向量广播除法 go = go / gl_block -// 封装 brcb gl → div 循环 → mask 尾部 → PIPE_BARRIER -// 用于 Stage2 段落C：go = go / gl_block（非TP1与TP1完全一致） -__aicore__ __attribute__((always_inline)) inline void PlatformDivVectorByBroadcast( - AscendC::LocalTensor go32_ubuf_tensor, - AscendC::LocalTensor tv32_ubuf_tensor, - AscendC::LocalTensor gl32_ubuf_tensor, - uint32_t sub_m, - uint32_t round_sub_m, - uint32_t round_v, - uint32_t __v, - uint32_t head_loop_idx) -{ - // gl_block = expand_to_block(gl), stored in tv - brcb_v(tv32_ubuf_tensor, - gl32_ubuf_tensor.ReinterpretCast()[head_loop_idx * 16], - 1, // dstBlockStride - 8, // dstRepeatStride - (uint8_t)(round_sub_m / FLOAT_BLOCK_SIZE) // repeat - ); - PIPE_BARRIER(V); - - // go = go / gl_block (full vector iterations) - SetVectorMask((uint64_t)-1, (uint64_t)-1); - for (uint32_t vdiv_idx = 0; vdiv_idx < __v / FLOAT_VECTOR_SIZE; ++vdiv_idx) { - div_v(go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[vdiv_idx * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor.template ReinterpretCast(), - (uint8_t)sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // dstRepeatStride - (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // src0RepeatStride - 1 // src1RepeatStride - ); - } - // tail mask iteration - if (__v % FLOAT_VECTOR_SIZE > 0) { - __set_mask(__v % FLOAT_VECTOR_SIZE); - div_v(go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - go32_ubuf_tensor[__v / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tv32_ubuf_tensor.template ReinterpretCast(), - (uint8_t)sub_m, // repeat - 1, // dstBlockStride - 1, // src0BlockStride - 0, // src1BlockStride - (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // dstRepeatStride - (uint8_t)(round_v / FLOAT_BLOCK_SIZE), // src0RepeatStride - 1 // src1RepeatStride - ); - SetVectorMask((uint64_t)-1, (uint64_t)-1); // fix hidden_size=96 - } - PIPE_BARRIER(V); -} - -// 平台函数：中间结果保存（段落D） -// 封装 SET_FLAG(V,MTE3,E5) → WAIT_FLAG → ub_to_gm go32 -// 用于 Stage2 head_loop>1 且非最后一轮时保存 go32 到 GM -// 三函数（非TP1/TP1/Tail）完全一致 -__aicore__ __attribute__((always_inline)) inline void PlatformIntermediateSave( - AscendC::GlobalTensor go_gm_tensor, - AscendC::LocalTensor go32_ubuf_tensor, - uint32_t sub_m, - uint32_t round_v) -{ - SET_FLAG(V, MTE3, EVENT_ID5); - WAIT_FLAG(V, MTE3, EVENT_ID5); - ub_to_gm( - go_gm_tensor, - go32_ubuf_tensor, - 0, - 1, - sub_m * round_v / FLOAT_BLOCK_SIZE, - 0, - 0 - ); -} - -// 平台函数：从 GM 搬入 go 到 UB（段落B 使用） -// 封装 gm_to_ub + SET_FLAG/WAIT_FLAG(MTE2, V, EVENT_ID0) -// 用于 Stage2MergeAccumulate head_loop>1 时加载上一轮 go -__aicore__ __attribute__((always_inline)) inline void PlatformLoadGoFromGM( - AscendC::GlobalTensor go_gm_tensor, - AscendC::LocalTensor go32_ubuf_tensor, - uint32_t sub_m, - uint32_t round_v) -{ - gm_to_ub( - go32_ubuf_tensor, go_gm_tensor, - 0, 1, sub_m * round_v / FLOAT_BLOCK_SIZE, 0, 0); - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); -} - -// 平台函数：向量加法 go = lo + go（段落B 使用） -// 封装 add_v + PIPE_BARRIER(V) -// 用于 Stage2MergeAccumulate 累加 lo 到 go -__aicore__ __attribute__((always_inline)) inline void PlatformAddLoToGo( - AscendC::LocalTensor go32_ubuf_tensor, - AscendC::LocalTensor lo_ubuf_tensor, - uint32_t sub_m, - uint32_t round_v) -{ - add_v(go32_ubuf_tensor, - go32_ubuf_tensor, lo_ubuf_tensor, - (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, 1, 1, 8, 8, 8); - PIPE_BARRIER(V); -} - -// 平台函数：FP32→FP16/BF16 类型转换 + V→MTE3 同步（段落C 使用） -// 封装 conv_v + SET_FLAG/WAIT_FLAG(V, MTE3, EVENT_ID0) -// 用于 Stage2FinalizeAndOutput go=castfp32to16(go) -__aicore__ __attribute__((always_inline)) inline void PlatformConvGoToOutput( - AscendC::LocalTensor go_ubuf_tensor, - AscendC::LocalTensor go32_ubuf_tensor, - uint32_t sub_m, - uint32_t round_v) -{ - conv_v(go_ubuf_tensor, - go32_ubuf_tensor, - (sub_m * round_v + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, // repeat - 1, 1, 4, 8); - SET_FLAG(V, MTE3, EVENT_ID0); - WAIT_FLAG(V, MTE3, EVENT_ID0); -} - -// 平台函数：DataCopyPad 输出到 GM（段落C 使用） -// 封装 head_res_row_num / numhead_per_process / tail_res_row_num 三段 DataCopyPad -// 用于 Stage2FinalizeAndOutput 输出 attention 结果到 o_gm -__aicore__ __attribute__((always_inline)) inline void PlatformDataCopyPadOutput( - AscendC::GlobalTensor o_gm_tensor, - AscendC::LocalTensor go_ubuf_tensor, - uint32_t sub_m, - uint32_t round_v, - uint32_t q_seq_len, - uint32_t numhead_per_process, - uint32_t head_res_row_num, - uint32_t head_start_sblock_idx, - uint32_t tail_res_row_num) -{ - uint32_t inner_o_gm_offset = 0; - uint32_t inner_go_ubuf_offset = 0; - - if (head_res_row_num != 0) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset + q_heads * __v * head_start_sblock_idx], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams(head_res_row_num, __v * 2, 0, - __v * (q_heads - 1) * 2, 0)); - inner_o_gm_offset += __v; - inner_go_ubuf_offset += head_res_row_num * __v; - } - - for (uint32_t i = 0; i < numhead_per_process; i++) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams(q_seq_len, __v * 2, 0, - __v * (q_heads - 1) * 2, 0)); - inner_o_gm_offset += __v; - inner_go_ubuf_offset += q_seq_len * __v; - } - - if (tail_res_row_num != 0) { - AscendC::DataCopyPad( - o_gm_tensor[inner_o_gm_offset], - go_ubuf_tensor[inner_go_ubuf_offset], - AscendC::DataCopyExtParams(tail_res_row_num, __v * 2, 0, - __v * (q_heads - 1) * 2, 0)); - } -} - -// 平台函数：Ring LSE copyout（段落C 使用） -// 封装 ln_v + PIPE_BARRIER + add_v + PIPE_BARRIER + conv_v + SET/WAIT_FLAG + ub_to_gm_align + SET/WAIT_FLAG -// 用于 Stage2FinalizeAndOutput IS_RING 路径的 LSE 输出 -__aicore__ __attribute__((always_inline)) inline void PlatformRingLSECopyout( - AscendC::LocalTensor lse32_ubuf_tensor, - AscendC::LocalTensor gl32_ubuf_tensor, - AscendC::LocalTensor gm32_ubuf_tensor, - AscendC::LocalTensor lse_conv_ubuf_tensor, - AscendC::GlobalTensor lse_gm_tensor, - uint32_t sub_m, - uint32_t head_loop, - uint64_t o_offset) -{ - uint32_t sub_m_d64 = (sub_m + 63) / 64; - uint32_t lenBurst = sizeof(OUT_DTYPE); - ln_v(lse32_ubuf_tensor, gl32_ubuf_tensor, - sub_m_d64, 1, 1, 8, 8); - PIPE_BARRIER(V); - add_v(lse32_ubuf_tensor, lse32_ubuf_tensor, - gm32_ubuf_tensor, sub_m_d64, 1, 1, 1, 8, 8, 8); - PIPE_BARRIER(V); - conv_v(lse_conv_ubuf_tensor, - lse32_ubuf_tensor, sub_m_d64, 1, 1, 4, 8); - SET_FLAG(V, MTE3, EVENT_ID1); - WAIT_FLAG(V, MTE3, EVENT_ID1); - ub_to_gm_align( - lse_gm_tensor[(int64_t)(o_offset / __k)], - lse_conv_ubuf_tensor, - 0, 1, lenBurst * sub_m * head_loop, 0, 0, 0, 0); - SET_FLAG(MTE3, V, EVENT_ID1); - WAIT_FLAG(MTE3, V, EVENT_ID1); -} - -// 平台函数：Tail copyout gl/gm/go（段落C Tail 使用） -// 封装 ub_to_gm(gl) + ub_to_gm(gm) + SET_FLAG/WAIT_FLAG(V,MTE3,E5) + ub_to_gm(go) -// 用于 TailStage2FinalizeAndOutput 输出 gl/gm/go 三个 GM buffer -__aicore__ __attribute__((always_inline)) inline void PlatformTailCopyout( - AscendC::GlobalTensor gl_gm_tensor, - AscendC::GlobalTensor gm_gm_tensor, - AscendC::GlobalTensor go_gm_tensor, - AscendC::LocalTensor gl32_ubuf_tensor, - AscendC::LocalTensor gm32_ubuf_tensor, - AscendC::LocalTensor go32_ubuf_tensor, - uint32_t sub_m, - uint32_t round_v) -{ - ub_to_gm( - gl_gm_tensor, gl32_ubuf_tensor, - 0, 1, sub_m / FLOAT_BLOCK_SIZE, 0, 0); - - ub_to_gm( - gm_gm_tensor, gm32_ubuf_tensor, - 0, 1, sub_m / FLOAT_BLOCK_SIZE, 0, 0); - - SET_FLAG(V, MTE3, EVENT_ID5); - WAIT_FLAG(V, MTE3, EVENT_ID5); - ub_to_gm( - go_gm_tensor, go32_ubuf_tensor, - 0, 1, sub_m * round_v / FLOAT_BLOCK_SIZE, 0, 0); -} - -// ==================== Stage1 / 通用数据搬运平台函数 ==================== - -// 平台函数：MTE2→V 通用同步（SET_FLAG + WAIT_FLAG, EVENT_ID0） -// 用于 LoadQKDataFP16: gm_to_ub 后通知 V 管道可读 -__aicore__ __attribute__((always_inline)) inline void PlatformMte2ToVSync() -{ - SET_FLAG(MTE2, V, EVENT_ID0); - WAIT_FLAG(MTE2, V, EVENT_ID0); -} - -// 平台函数：V→MTE3 同步（SET_FLAG + WAIT_FLAG, EVENT_ID0） -// 用于 QuantizeAndOutput: conv_v 后通知 MTE3 可输出 -__aicore__ __attribute__((always_inline)) inline void PlatformVToMte3Sync() -{ - SET_FLAG(V, MTE3, EVENT_ID0); - WAIT_FLAG(V, MTE3, EVENT_ID0); -} - -// 平台函数：V→MTE2 通知（SET_FLAG + PIPE_BARRIER） -// 用于 QuantizeAndOutput: ReduceSum 后通知 MTE2 可搬入下一轮 -__aicore__ __attribute__((always_inline)) inline void PlatformVToMte2Notify() -{ - SET_FLAG(V, MTE2, EVENT_ID0); - PIPE_BARRIER(V); -} - -// 平台函数：V→MTE2 等待（WAIT_FLAG, EVENT_ID0） -// 用于 HeadLoop: 等待上一轮 V→MTE2 通知 -__aicore__ __attribute__((always_inline)) inline void PlatformVToMte2WaitEvent0() -{ - WAIT_FLAG(V, MTE2, EVENT_ID0); -} - -// 平台函数：MTE3→MTE2 等待（WAIT_FLAG, EVENT_ID4） -// 用于 HeadLoop: 等待 MTE3→MTE2 管道就绪 -__aicore__ __attribute__((always_inline)) inline void PlatformMte3ToMte2WaitEvent4() -{ - WAIT_FLAG(MTE3, MTE2, EVENT_ID4); -} - -// 平台函数：V→MTE2 通知（SET_FLAG, EVENT_ID0）— 仅 SET -// 用于 HeadLoop: 通知 MTE2 可读 go -__aicore__ __attribute__((always_inline)) inline void PlatformVToMte2SetEvent0() -{ - SET_FLAG(V, MTE2, EVENT_ID0); -} - -// 平台函数：MTE3→MTE2 通知（SET_FLAG, EVENT_ID4）— 仅 SET -// 用于 HeadLoop 尾部: 通知 MTE2 下一轮可搬入 -__aicore__ __attribute__((always_inline)) inline void PlatformMte3ToMte2SetEvent4() -{ - SET_FLAG(MTE3, MTE2, EVENT_ID4); -} - -// 平台函数：MTE3→MTE2 等待（WAIT_FLAG, EVENT_ID1）— TP1 专用 -// 用于 TP1 HeadLoop: 等待 MTE3→MTE2 EVENT_ID1 -__aicore__ __attribute__((always_inline)) inline void PlatformMte3ToMte2WaitEvent1() -{ - WAIT_FLAG(MTE3, MTE2, EVENT_ID1); -} - -// 平台函数：MTE3→MTE2 通知（SET_FLAG, EVENT_ID1）— TP1 专用 -// 用于 TP1 HeadLoop 尾部: 通知 MTE2 EVENT_ID1 -__aicore__ __attribute__((always_inline)) inline void PlatformMte3ToMte2SetEvent1() -{ - SET_FLAG(MTE3, MTE2, EVENT_ID1); -} - -// 平台函数：PIPE_BARRIER(ALL) -// 用于 Tail HeadLoop 尾部全管道屏障 -__aicore__ __attribute__((always_inline)) inline void PlatformPipeBarrierAll() -{ - PIPE_BARRIER(ALL); -} - -// ==================== 通用数据搬运平台函数 ==================== - -// 平台函数：gm_to_ub（通用 GM→UB 搬入，float 类型） -__aicore__ __attribute__((always_inline)) inline void PlatformGmToUbFloat( - AscendC::LocalTensor dst, - AscendC::GlobalTensor src, - uint32_t lenBurst) -{ - gm_to_ub(dst, src, 0, 1, lenBurst, 0, 0); -} - -// 平台函数：gm_to_ub（通用 GM→UB 搬入，模板类型） -template -__aicore__ __attribute__((always_inline)) inline void PlatformGmToUb( - AscendC::LocalTensor dst, - AscendC::GlobalTensor src, - uint32_t lenBurst) -{ - gm_to_ub(dst, src, 0, 1, lenBurst, 0, 0); -} - -// 平台函数：gm_to_ub（INT32 类型，CeilDiv 对齐） -__aicore__ __attribute__((always_inline)) inline void PlatformGmToUbInt32( - AscendC::LocalTensor dst, - AscendC::GlobalTensor src, - uint32_t sub_m, uint32_t qk_round_n) -{ - gm_to_ub( - dst, src, 0, 1, CeilDiv(sub_m * qk_round_n), 0, 0); -} - -// 平台函数：gm_to_ub_align（对齐 GM→UB，mmScaleType） -__aicore__ __attribute__((always_inline)) inline void PlatformGmToUbAlignMmScale( - AscendC::LocalTensor dst, - AscendC::GlobalTensor src, - uint32_t sub_m) -{ - gm_to_ub_align( - dst, src, 0, 1, sub_m * sizeof(mmScaleType), 0, 0, 0, 0); -} - -// 平台函数：ub_to_gm（通用 UB→GM 搬出，模板类型） -template -__aicore__ __attribute__((always_inline)) inline void PlatformUbToGm( - AscendC::GlobalTensor dst, - AscendC::LocalTensor src, - uint32_t lenBurst) -{ - ub_to_gm(dst, src, 0, 1, lenBurst, 0, 0); -} - -// 平台函数：ub_to_ub（通用 UB→UB 搬移，float 类型） -__aicore__ __attribute__((always_inline)) inline void PlatformUbToUb( - AscendC::LocalTensor dst, - AscendC::LocalTensor src, - uint32_t lenBurst) -{ - ub_to_ub(dst, src, 0, 1, lenBurst, 0, 0); -} - -// 平台函数：ub_to_ub（UB→UB，多 nBurst） -__aicore__ __attribute__((always_inline)) inline void PlatformUbToUbMulti( - AscendC::LocalTensor dst, - AscendC::LocalTensor src, - uint32_t nBurst, - uint32_t lenBurst, - uint32_t srcGap, - uint32_t dstGap) -{ - ub_to_ub(dst, src, 0, nBurst, lenBurst, srcGap, dstGap); -} - -// 平台函数：DataCopyPad（不规则 mask 加载） -__aicore__ __attribute__((always_inline)) inline void PlatformDataCopyPadMask( - AscendC::LocalTensor dst, - AscendC::GlobalTensor src, - uint32_t cur_q_seqlen, - uint32_t qk_n, - uint32_t maxKVSeqLen, - uint32_t qk_round_n) -{ - uint32_t aligned_mask_copy_len = RoundUp(qk_n); - uint32_t mask_dst_stride = (qk_round_n - aligned_mask_copy_len) / BLOCK_SIZE; - AscendC::DataCopyPad( - dst, src, - AscendC::DataCopyExtParams(cur_q_seqlen, qk_n * 2, maxKVSeqLen * 2 - qk_n * 2, mask_dst_stride, 0), - AscendC::DataCopyPadExtParams(false, 0, 0, 0) - ); -} - -// 平台函数：DataCopy（规则 mask 加载） -__aicore__ __attribute__((always_inline)) inline void PlatformDataCopyMask( - AscendC::LocalTensor dst, - AscendC::GlobalTensor src, - uint32_t cur_q_seqlen, - uint32_t qk_round_n, - uint32_t maxKVSeqLen) -{ - AscendC::DataCopy( - dst, src, - AscendC::DataCopyParams( - cur_q_seqlen, - qk_round_n * 2 / 32, - MASK_COLUMNS * 2 / 32 - qk_round_n * 2 / 32, - 0) - ); -} - -// 平台函数：Cast mask→float -__aicore__ __attribute__((always_inline)) inline void PlatformCastMaskToFloat( - AscendC::LocalTensor dst, - AscendC::LocalTensor src, - uint32_t count) -{ - AscendC::Cast(dst, src, AscendC::RoundMode::CAST_NONE, count); -} - -// 平台函数：SetVectorMask 全置 -1 -__aicore__ __attribute__((always_inline)) inline void PlatformSetVectorMaskFull() -{ - SetVectorMask((uint64_t)-1, (uint64_t)-1); -} - -// 平台函数：__set_mask 设置部分 mask -__aicore__ __attribute__((always_inline)) inline void PlatformSetMask(uint32_t mask) -{ - __set_mask(mask); -} - -// ==================== 向量计算平台函数 ==================== - -// 平台函数：muls_v（标量乘法，循环 + tail mask） -__aicore__ __attribute__((always_inline)) inline void PlatformMulsVLoop( - AscendC::LocalTensor dst, - AscendC::LocalTensor src, - float scalar, - uint32_t sub_m, - uint32_t qk_n, - uint32_t qk_round_n) -{ - for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { - muls_v( - dst[vadd_idx * FLOAT_VECTOR_SIZE], - src[vadd_idx * FLOAT_VECTOR_SIZE], - scalar, sub_m, 1, 1, - qk_round_n / FLOAT_BLOCK_SIZE, - qk_round_n / FLOAT_BLOCK_SIZE); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - muls_v( - dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - scalar, sub_m, 1, 1, - qk_round_n / FLOAT_BLOCK_SIZE, - qk_round_n / FLOAT_BLOCK_SIZE); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); -} - -// 平台函数：Add（AscendC::Add，FP32 逐元素加） -__aicore__ __attribute__((always_inline)) inline void PlatformAddFloat( - AscendC::LocalTensor dst, - AscendC::LocalTensor src0, - AscendC::LocalTensor src1, - uint32_t count) -{ - AscendC::Add(dst, src0, src1, count); -} - -// 平台函数：max_v（逐元素取最大值） -__aicore__ __attribute__((always_inline)) inline void PlatformMaxV( - AscendC::LocalTensor dst, - AscendC::LocalTensor src0, - AscendC::LocalTensor src1, - uint32_t repeat) -{ - max_v( - dst, src0, src1, (uint8_t)repeat, 1, 1, 1, 8, 8, 8); -} - -// 平台函数：sub_v（逐元素减法） -__aicore__ __attribute__((always_inline)) inline void PlatformSubV( - AscendC::LocalTensor dst, - AscendC::LocalTensor src0, - AscendC::LocalTensor src1, - uint32_t repeat) -{ - sub_v( - dst, src0, src1, (uint8_t)repeat, 1, 1, 1, 8, 8, 8); -} - -// 平台函数：exp_v（指数运算） -__aicore__ __attribute__((always_inline)) inline void PlatformExpV( - AscendC::LocalTensor dst, - AscendC::LocalTensor src, - uint32_t repeat) -{ - exp_v( - dst, src, (uint8_t)repeat, 1, 1, 8, 8); -} - -// 平台函数：mul_v（逐元素乘法） -__aicore__ __attribute__((always_inline)) inline void PlatformMulV( - AscendC::LocalTensor dst, - AscendC::LocalTensor src0, - AscendC::LocalTensor src1, - uint32_t repeat) -{ - mul_v( - dst, src0, src1, (uint8_t)repeat, 1, 1, 1, 8, 8, 8); -} - -// 平台函数：add_v（逐元素加法） -__aicore__ __attribute__((always_inline)) inline void PlatformAddV( - AscendC::LocalTensor dst, - AscendC::LocalTensor src0, - AscendC::LocalTensor src1, - uint32_t repeat) -{ - add_v( - dst, src0, src1, (uint8_t)repeat, 1, 1, 1, 8, 8, 8); -} - -// 平台函数：muls_v（标量乘法，单次调用） -__aicore__ __attribute__((always_inline)) inline void PlatformMulsV( - AscendC::LocalTensor dst, - AscendC::LocalTensor src, - float scalar, - uint32_t repeat) -{ - muls_v( - dst, src, scalar, (uint8_t)repeat, 1, 1, 8, 8); -} - -// 平台函数：brcb_v（广播，uint32_t） -__aicore__ __attribute__((always_inline)) inline void PlatformBrcbV( - AscendC::LocalTensor dst, - AscendC::LocalTensor src, - uint32_t repeat) -{ - brcb_v(dst, src, 1, 8, (uint8_t)repeat); -} - -// 平台函数：brcb_v（广播，float） -__aicore__ __attribute__((always_inline)) inline void PlatformBrcbVFloat( - AscendC::LocalTensor dst, - AscendC::LocalTensor src, - uint32_t round_sub_m) -{ - brcb_v( - dst.ReinterpretCast(), - src.ReinterpretCast(), - 1, 8, (uint8_t)(round_sub_m / FLOAT_BLOCK_SIZE)); -} - -// 平台函数：conv_v（float→OUT_DTYPE 转换） -__aicore__ __attribute__((always_inline)) inline void PlatformConvVToOutput( - AscendC::LocalTensor dst, - AscendC::LocalTensor src, - uint32_t repeat) -{ - conv_v( - dst, src, (uint8_t)repeat, 1, 1, 4, 8); -} - -// 平台函数：conv_v（int32_t→float 转换，含 repeat_times<255 分支） -__aicore__ __attribute__((always_inline)) inline void PlatformConvInt32ToFloat( - AscendC::LocalTensor dst, - AscendC::LocalTensor temp, - uint32_t count) -{ - uint32_t repeat_times = (count + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE; - if (repeat_times < 255) { - conv_v( - dst, temp, (uint8_t)repeat_times, 1, 1, 8, 8); - } else { - for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { - conv_v( - dst[vconv_idx * count / 2], temp[vconv_idx * count / 2], - (uint8_t)((count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE), - 1, 1, 8, 8); - } - } -} - -// 平台函数：conv_v（float→half 转换，含 repeat_times<255 分支） -__aicore__ __attribute__((always_inline)) inline void PlatformConvFloatToHalf( - AscendC::LocalTensor dst, - uint32_t count) -{ - uint32_t repeat_times = (count + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE; - if (repeat_times < 255) { - conv_v( - dst.template ReinterpretCast(), - dst.template ReinterpretCast(), - (uint8_t)repeat_times, 1, 1, 4, 8); - } else { - for (uint64_t vconv_idx = 0; vconv_idx < 2; ++vconv_idx) { - conv_v( - dst.template ReinterpretCast()[vconv_idx * count / 2], - dst.template ReinterpretCast()[vconv_idx * count / 2], - (uint8_t)((count / 2 + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE), - 1, 1, 4, 8); - } - } - PIPE_BARRIER(V); -} - -// 平台函数：Cast 逐行转换（循环 + tail mask） -__aicore__ __attribute__((always_inline)) inline void PlatformCastHalfToInt8( - AscendC::LocalTensor dst, - uint32_t qk_n, - uint32_t qk_round_n, - uint32_t sub_m) -{ - for (uint32_t row_idx = 0; row_idx < qk_n / HALF_VECTOR_SIZE; ++row_idx) { - AscendC::Cast( - dst.template ReinterpretCast()[row_idx * HALF_VECTOR_SIZE], - dst.template ReinterpretCast()[row_idx * HALF_VECTOR_SIZE], - AscendC::RoundMode::CAST_RINT, - (uint64_t)0, sub_m, - {1, 1, (uint8_t)((qk_round_n) / BLOCK_SIZE), (uint8_t)(qk_round_n / BLOCK_SIZE)}); - } - if (qk_n % HALF_VECTOR_SIZE > 0) { - __set_mask(qk_n % HALF_VECTOR_SIZE); - AscendC::Cast( - dst.template ReinterpretCast()[qk_n / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], - dst.template ReinterpretCast()[qk_n / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], - AscendC::RoundMode::CAST_RINT, - (uint64_t)0, sub_m, - {1, 1, (uint8_t)((qk_round_n) / BLOCK_SIZE), (uint8_t)(qk_round_n / BLOCK_SIZE)}); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); -} - -// ==================== Tensor/Reduce 系列平台函数 ==================== -// 以下函数从 bs 层整体搬入，它们是平台原语+循环的封装，属于平台操作 - -// 平台函数：Tensor 逐行减法（brcb + sub_v 循环 + tail mask） -__aicore__ __attribute__((always_inline)) inline void PlatformTensorSubValueRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& MaxTensor, - const AscendC::LocalTensor& tempMaxTensor, - uint32_t sub_m, - uint32_t round_sub_m, - uint32_t qk_n, - uint32_t qk_round_n) -{ - brcb_v( - tempMaxTensor.ReinterpretCast(), - MaxTensor.ReinterpretCast(), - 1, 8, (uint8_t)(round_sub_m / FLOAT_BLOCK_SIZE)); - PIPE_BARRIER(V); - for (uint32_t sub_v_idx = 0; sub_v_idx < qk_n / FLOAT_VECTOR_SIZE; ++sub_v_idx) { - sub_v( - dst[sub_v_idx * FLOAT_VECTOR_SIZE], - src[sub_v_idx * FLOAT_VECTOR_SIZE], - tempMaxTensor, (uint8_t)sub_m, 1, 1, 0, - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), 1); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - sub_v( - dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - tempMaxTensor, (uint8_t)sub_m, 1, 1, 0, - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), 1); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); -} - -// 平台函数：Tensor 逐行除法（div_v 循环 + tail mask） -__aicore__ __attribute__((always_inline)) inline void PlatformTensorDivRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& src1, - uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n) -{ - PIPE_BARRIER(V); - for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { - div_v( - dst[vadd_idx * FLOAT_VECTOR_SIZE], - src[vadd_idx * FLOAT_VECTOR_SIZE], - src1, (uint8_t)sub_m, 1, 1, 0, - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), 1); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - div_v( - dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src1, (uint8_t)sub_m, 1, 1, 0, - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), 1); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); -} - -// 平台函数：Tensor 行归约求最大值（cmax_v + max_v 循环 + tail mask） -__aicore__ __attribute__((always_inline)) inline void PlatformReduceMaxRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& tempTensor, - uint32_t sub_m, - uint32_t qk_n, - uint32_t qk_round_n) -{ - if (qk_n <= FLOAT_VECTOR_SIZE) { - __set_mask(qk_n); - cmax_v( - dst, src, (uint8_t)sub_m, 1, 1, (uint16_t)(qk_round_n / FLOAT_BLOCK_SIZE)); - } else { - ub_to_ub( - tempTensor, src, 0, sub_m, - HALF_VECTOR_SIZE / BLOCK_SIZE, - (qk_round_n - FLOAT_VECTOR_SIZE) / FLOAT_BLOCK_SIZE, 0); - PIPE_BARRIER(V); - for (uint32_t rowmax_idx = 1; rowmax_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowmax_idx) { - max_v( - tempTensor, tempTensor, src[rowmax_idx * FLOAT_VECTOR_SIZE], - (uint8_t)sub_m, 1, 1, 1, 8, 8, (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE)); - PIPE_BARRIER(V); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - max_v( - tempTensor, tempTensor, - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - (uint8_t)sub_m, 1, 1, 1, 8, 8, (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE)); - } - PIPE_BARRIER(V); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - cmax_v( - dst, tempTensor, (uint8_t)sub_m, 1, 1, 8); - } - SetVectorMask((uint64_t)-1, (uint64_t)-1); - PIPE_BARRIER(V); -} - -// 平台函数：Tensor 行归约求和（cadd_v + add_v 循环 + tail mask） -__aicore__ __attribute__((always_inline)) inline void PlatformReduceSumRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - uint32_t sub_m, - uint32_t qk_n, - uint32_t qk_round_n) -{ - if (qk_n <= FLOAT_VECTOR_SIZE) { - __set_mask(qk_n); - cadd_v( - dst, src, (uint8_t)sub_m, 1, 1, (uint16_t)(qk_round_n / FLOAT_BLOCK_SIZE)); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } else { - for (uint32_t rowsum_idx = 1; rowsum_idx < qk_n / FLOAT_VECTOR_SIZE; ++rowsum_idx) { - add_v( - src, src, src[rowsum_idx * FLOAT_VECTOR_SIZE], - (uint8_t)sub_m, 1, 1, 1, - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE)); - PIPE_BARRIER(V); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - add_v( - src, src, - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - (uint8_t)sub_m, 1, 1, 1, - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE)); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); - cadd_v( - dst, src, (uint8_t)sub_m, 1, 1, (uint16_t)(qk_round_n / FLOAT_BLOCK_SIZE)); - } -} - -// 平台函数：Tensor 逐行乘法（mul_v 循环 + tail mask） -__aicore__ __attribute__((always_inline)) inline void PlatformTensorMulRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& src1, - uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t src1BlockStride) -{ - PIPE_BARRIER(V); - for (uint32_t vadd_idx = 0; vadd_idx < qk_n / FLOAT_VECTOR_SIZE; ++vadd_idx) { - mul_v( - dst[vadd_idx * FLOAT_VECTOR_SIZE], - src[vadd_idx * FLOAT_VECTOR_SIZE], - src1, (uint8_t)sub_m, 1, 1, (uint8_t)src1BlockStride, - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), 1); - } - if (qk_n % FLOAT_VECTOR_SIZE > 0) { - __set_mask(qk_n % FLOAT_VECTOR_SIZE); - mul_v( - dst[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src[qk_n / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], - src1, (uint8_t)sub_m, 1, 1, (uint8_t)src1BlockStride, - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), - (uint8_t)(qk_round_n / FLOAT_BLOCK_SIZE), 1); - SetVectorMask((uint64_t)-1, (uint64_t)-1); - } - PIPE_BARRIER(V); -} - -// ==================== DeQuant/Quant 系列平台函数 ==================== - -// 平台函数：加载 deScale + online 乘 quantScale -__aicore__ __attribute__((always_inline)) inline void PlatformLoadDeScaleAndOnlineMul( - const AscendC::GlobalTensor& deScaleGm, - AscendC::LocalTensor deScaleUb, - AscendC::LocalTensor quantScale, - uint32_t sub_m, bool online) -{ - gm_to_ub_align( - deScaleUb, deScaleGm, 0, 1, sub_m * sizeof(mmScaleType), 0, 0, 0, 0); - if (online) { - PlatformMte2ToVSyncEvent2(); - PlatformTensorMulRepeatM(deScaleUb, deScaleUb, quantScale, 1, sub_m, RoundUp<16>(sub_m), 1); - } -} - -// 平台函数：加载 src(int32) + brcb 广播 deScale → tempScale -__aicore__ __attribute__((always_inline)) inline void PlatformLoadSrcAndBrcbScale( - const AscendC::GlobalTensor& src, - AscendC::LocalTensor temp, - AscendC::LocalTensor deScaleUb, - AscendC::LocalTensor tempScale, - uint32_t sub_m, uint32_t qk_round_n, bool move_tensor) -{ - if (move_tensor) { - gm_to_ub( - temp, src, 0, 1, CeilDiv(sub_m * qk_round_n), 0, 0); - } - PlatformMte2ToVSyncEvent0(); - brcb_v( - tempScale.template ReinterpretCast(), - deScaleUb.template ReinterpretCast(), - 1, 8, (uint8_t)(RoundUp<16>(sub_m) / FLOAT_BLOCK_SIZE)); - PlatformVPipeBarrier(); -} - -// 平台函数：INT32→FP32 转换(conv_v) + 乘 tempScale -__aicore__ __attribute__((always_inline)) inline void PlatformConvInt32ToFP32AndMul( - AscendC::LocalTensor dst, - AscendC::LocalTensor temp, - AscendC::LocalTensor tempScale, - uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n) -{ - uint32_t count = sub_m * qk_round_n; - PlatformConvInt32ToFloat(dst, temp, count); - PlatformTensorMulRepeatM(dst, dst, tempScale, sub_m, qk_n, qk_round_n, 0); - PlatformVPipeBarrier(); -} - -// 平台函数：逐 Head 反量化（编排：加载+转换+乘法） -__aicore__ __attribute__((always_inline)) inline void PlatformDeQuantPerHead( - const AscendC::GlobalTensor& deScaleGm, - const AscendC::GlobalTensor& src, - AscendC::LocalTensor dst, - AscendC::LocalTensor temp, - AscendC::LocalTensor deScaleUb, - AscendC::LocalTensor tempScale, - AscendC::LocalTensor quantScale, - uint32_t sub_m, - uint32_t qk_n, - uint32_t qk_round_n, - bool online, - bool move_tensor) -{ - PlatformLoadDeScaleAndOnlineMul(deScaleGm, deScaleUb, quantScale, sub_m, online); - PlatformLoadSrcAndBrcbScale(src, temp, deScaleUb, tempScale, sub_m, qk_round_n, move_tensor); - PlatformConvInt32ToFP32AndMul(dst, temp, tempScale, sub_m, qk_n, qk_round_n); -} - -// 平台函数：逐 Token 量化（FP32 → FP16 → INT8） -__aicore__ __attribute__((always_inline)) inline void PlatformQuantPerToken( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& scale, - uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t pQuantOnline) -{ - if (pQuantOnline) { - PlatformTensorDivRepeatM(dst.template ReinterpretCast(), src, scale, sub_m, qk_n, qk_round_n); - } else { - PlatformTensorMulRepeatM(dst.template ReinterpretCast(), src, scale, sub_m, qk_n, qk_round_n, 0); - } - uint32_t count = sub_m * qk_round_n; - PlatformConvFloatToHalf(dst, count); - PlatformCastHalfToInt8(dst, qk_n, qk_round_n, sub_m); -} \ No newline at end of file diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h deleted file mode 100644 index 2240fc5..0000000 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_aiv_bs.h +++ /dev/null @@ -1,1303 +0,0 @@ -// Copyright 2025 The xLLM Authors. 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: -// -// https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE -// -// ==================== AIV 业务逻辑层（第二层）==================== -// 本文件仅被 MLADecoderAiv 类 include，包含 Vector 侧业务调度函数。 -// AIC 侧业务函数在 multi_latent_attention_bs.h 中。 - -#pragma once - -// ====== AIV refactor: context struct + init functions ====== - -// 非 TP1 路径上下文 -struct VectorContext { - // input params - uint32_t cur_batch; - uint32_t start_head; - uint32_t cur_nIndx; - uint32_t cur_q_seqlen; - uint32_t cur_kv_seqlen; - uint32_t cur_head_num; - uint32_t offset_tiling; - - // addresses - uint64_t addr_o_scalar; - uint64_t addr_mask_scalar; - uint32_t mask_offset; - - // loop & size - uint32_t pp_n_scalar; - uint32_t sub_n_loop; - uint32_t real_n_loop; - uint32_t n_loop; - - // QK dims - uint32_t qk_n; - uint32_t qk_round_n; - uint32_t qk_n_2; - uint32_t qk_round_n_2; - - // head split - uint32_t sub_head_num; - uint32_t sub_m; - uint32_t head_idx; - uint64_t o_offset; - - // tail info - uint32_t tail_len; - bool prev_tail_mask; -}; - -__aicore__ __attribute__((always_inline)) inline void InitVectorContext( - VectorContext &ctx, uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, - uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, - uint32_t offset_tiling) -{ - ctx.cur_batch = cur_batch; - ctx.start_head = start_head; - ctx.cur_nIndx = cur_nIndx; - ctx.cur_q_seqlen = cur_q_seqlen; - ctx.cur_kv_seqlen = cur_kv_seqlen; - ctx.cur_head_num = cur_head_num; - ctx.offset_tiling = offset_tiling; - - // addr - uint32_t addr_o_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 4 + offset_tiling)); - uint32_t addr_o_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 5 + offset_tiling)); - ctx.addr_o_scalar = (uint64_t)(((uint64_t)addr_o_high32) << 32 | addr_o_loww32); - - uint32_t addr_mask_high32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 6 + offset_tiling)); - uint32_t addr_mask_loww32 = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 7 + offset_tiling)); - ctx.addr_mask_scalar = (uint64_t)(((uint64_t)addr_mask_high32) << 32 | addr_mask_loww32); - ctx.mask_offset = ctx.addr_mask_scalar; - - // loop & size - ctx.pp_n_scalar = block_size; - ctx.sub_n_loop = ctx.pp_n_scalar / block_size; - ctx.real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; - ctx.n_loop = (cur_kv_seqlen + ctx.pp_n_scalar - 1) / ctx.pp_n_scalar; - - // QK dims - ctx.qk_n = ctx.pp_n_scalar; - ctx.qk_round_n = RoundUp(ctx.qk_n); - ctx.qk_n_2 = ctx.pp_n_scalar; - ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); - - // head split - ctx.sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; - ctx.sub_m = ctx.sub_head_num * cur_q_seqlen; - ctx.head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; - ctx.o_offset = ctx.addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; - - // tail info - ctx.tail_len = cur_kv_seqlen - (ctx.n_loop - 1) * ctx.pp_n_scalar; - ctx.prev_tail_mask = (ctx.n_loop > 1 && ctx.tail_len < cur_q_seqlen - 1); -} - -// TP1 路径上下文 -struct VectorTP1Context { - // input params - uint32_t cur_batch; - uint32_t start_head; - uint32_t cur_nIndx; - uint32_t cur_q_seqlen; - uint32_t cur_kv_seqlen; - uint32_t cur_head_num; - uint32_t offset_tiling; - - // addresses - uint64_t addr_o_scalar; - uint64_t addr_mask_scalar; - uint32_t mask_offset; - - // loop & size - uint32_t pp_n_scalar; - uint32_t sub_n_loop; - uint32_t real_n_loop; - uint32_t n_loop; - - // QK dims - uint32_t qk_n; - uint32_t qk_round_n; - uint32_t qk_n_2; - uint32_t qk_round_n_2; - - // head split - uint32_t sub_head_num; - uint32_t sub_m; - uint32_t head_idx; - uint64_t o_offset; - - // TP1 specific - uint32_t s_block_stack; - uint32_t m_slice; - uint32_t m_end; -}; - -__aicore__ __attribute__((always_inline)) inline void InitVectorTP1Context( - VectorTP1Context &ctx, uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, - uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, - uint32_t offset_tiling) -{ - ctx.cur_batch = cur_batch; - ctx.start_head = start_head; - ctx.cur_nIndx = cur_nIndx; - ctx.cur_q_seqlen = cur_q_seqlen; - ctx.cur_kv_seqlen = cur_kv_seqlen; - ctx.cur_head_num = cur_head_num; - ctx.offset_tiling = offset_tiling; - - // addr - uint32_t prev_task = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); - ctx.addr_o_scalar = prev_task * q_heads * embedding_size; - ctx.addr_mask_scalar = 0; - ctx.mask_offset = 0; - - // loop & size - ctx.pp_n_scalar = block_size; - ctx.sub_n_loop = ctx.pp_n_scalar / block_size; - ctx.real_n_loop = (cur_kv_seqlen + block_size - 1) / block_size; - ctx.n_loop = (cur_kv_seqlen + ctx.pp_n_scalar - 1) / ctx.pp_n_scalar; - - // QK dims - ctx.qk_n = ctx.pp_n_scalar; - ctx.qk_round_n = RoundUp(ctx.qk_n); - ctx.qk_n_2 = ctx.pp_n_scalar; - ctx.qk_round_n_2 = RoundUp(ctx.qk_n_2); - - // head split - ctx.sub_head_num = (sub_block_idx == 1) ? (cur_head_num - cur_head_num / 2) : cur_head_num / 2; - ctx.sub_m = ctx.sub_head_num * cur_q_seqlen; - ctx.head_idx = (sub_block_idx == 0) ? start_head : start_head + cur_head_num / 2 * cur_q_seqlen; - ctx.o_offset = ctx.addr_o_scalar + start_head * embedding_size + sub_block_idx * cur_head_num / 2 * embedding_size; - - // TP1 specific - ctx.s_block_stack = 4; - ctx.m_slice = FLOAT_VECTOR_SIZE / ctx.s_block_stack; - ctx.m_end = (ctx.sub_m + ctx.m_slice - 1) / ctx.m_slice; -} - -// ====== 非 TP1 路径：Phase 1 / Phase 2 业务调度 ====== - -// 业务子函数1a：QK 数据加载（INT8 路径） -// INT8 量化 QK 反量化 + RoPE 残差加载 + Add 融合 -// 平台调用：PlatformVToMte2Wait / PlatformMte2ToVSync / PlatformVPipeBarrier -__aicore__ __attribute__((always_inline)) inline void LoadQKDataInt8( - AscendC::GlobalTensor s_gm_tensor, - AscendC::GlobalTensor s_rope_gm_tensor, - uint32_t qk_n, - uint32_t qk_round_n, - uint32_t sub_m, - const uint32_t head_idx) -{ - PlatformVToMte2Wait(); - PlatformDeQuantPerHead( - deq_scale_gm_tensor_q1[head_idx], s_gm_tensor, - ls32_quant_ubuf_tensor, ls32_quant_ubuf_tensor.template ReinterpretCast(), - descale_q1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, qk_n, qk_round_n, 0, 1); - PlatformGmToUbFloat( - ls32_ubuf_tensor.template ReinterpretCast(), - s_rope_gm_tensor, - sub_m * qk_round_n / FLOAT_BLOCK_SIZE); - PlatformMte2ToVSync(); - PlatformAddFloat(ls32_ubuf_tensor, ls32_ubuf_tensor, ls32_quant_ubuf_tensor, sub_m * qk_round_n); - PlatformVPipeBarrier(); -} - -// 业务子函数1b：QK 数据加载（非INT8 路径） -// FP16/BF16 QK 直接加载 + 可选 mask 加载(DataCopyPad/DataCopy) + Cast -// 平台调用：PlatformVToMte2Wait / PlatformMte2ToVSync -__aicore__ __attribute__((always_inline)) inline void LoadQKDataFP16( - AscendC::GlobalTensor s_gm_tensor, - AscendC::GlobalTensor mask_gm_tensor, - uint32_t qk_n, - uint32_t qk_round_n, - uint32_t sub_m, - uint32_t cur_q_seqlen, - bool need_mask) -{ - PlatformVToMte2Wait(); - PlatformGmToUb( - ls32_ubuf_tensor.template ReinterpretCast(), - s_gm_tensor, - sub_m * qk_round_n / FLOAT_BLOCK_SIZE); - - // mask 加载：mask_type==3 不规则 padding, mask_type==4 规则, 默认不加载 - if (mask_type == 3) { - PlatformDataCopyPadMask(mask_ubuf_tensor, mask_gm_tensor, - cur_q_seqlen, qk_n, maxKVSeqLen, qk_round_n); - } else if (need_mask && mask_type == 4) { - PlatformDataCopyMask(mask_ubuf_tensor, mask_gm_tensor, - cur_q_seqlen, qk_round_n, maxKVSeqLen); - } - - PlatformMte2ToVSync(); - - if (mask_type == 3 || (need_mask && mask_type == 4)) { - PlatformCastMaskToFloat(mask32_ubuf_tensor, mask_ubuf_tensor, - cur_q_seqlen * qk_round_n); - } -} - -// 业务子函数2：QK 缩放 + 非INT8 mask Add -// ls = ls × scale(tor)，非INT8 路径再叠加 mask -__aicore__ __attribute__((always_inline)) inline void ScaleAndMask( - uint32_t qk_n, - uint32_t qk_round_n, - uint32_t sub_m, - uint32_t cur_q_seqlen, - bool need_mask -) -{ - PlatformMulsVLoop(ls32_ubuf_tensor, ls32_ubuf_tensor, tor, sub_m, qk_n, qk_round_n); - - if constexpr (tilingKeyType != TilingKeyType::TILING_INT8_DATA) { - if (mask_type == 3 || (need_mask && mask_type == 4)) { - uint32_t cur_compute_head_num = sub_m / cur_q_seqlen; - for (uint32_t i = 0; i < cur_compute_head_num; i++) { - PlatformAddFloat( - ls32_ubuf_tensor[cur_q_seqlen * qk_round_n * i], - ls32_ubuf_tensor[cur_q_seqlen * qk_round_n * i], - mask32_ubuf_tensor, - cur_q_seqlen * qk_round_n); - } - PlatformVPipeBarrier(); - } - } -} - -// 业务子函数3：Online Softmax 状态更新 -// lm = rowmax(ls) → hm = max(lm, gm) → dm = gm - hm → gm = hm -__aicore__ __attribute__((always_inline)) inline void UpdateSoftmaxState( - AscendC::LocalTensor dm32_ubuf_tensor, - uint32_t n_idx, - uint32_t sub_m, - uint32_t qk_n, - uint32_t qk_round_n -) -{ - uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - - // *** lm = rowmax(ls) - PlatformReduceMaxRepeatM(lm32_ubuf_tensor, ls32_ubuf_tensor, lp32_ubuf_tensor, sub_m, qk_n, qk_round_n); - if (n_idx != 0) { - // *** hm = vmax(lm, gm) - PlatformMaxV(hm32_ubuf_tensor, lm32_ubuf_tensor, gm32_ubuf_tensor, sub_m_d64); - PlatformVPipeBarrier(); - // *** dm = gm - hm - PlatformSubV(dm32_ubuf_tensor, gm32_ubuf_tensor, hm32_ubuf_tensor, sub_m_d64); - PlatformVPipeBarrier(); - } else { - // *** hm = lm - PlatformUbToUb(hm32_ubuf_tensor, lm32_ubuf_tensor, round_sub_m / FLOAT_BLOCK_SIZE); - PlatformVPipeBarrier(); - } - // *** gm = hm - PlatformUbToUb(gm32_ubuf_tensor, hm32_ubuf_tensor, round_sub_m / FLOAT_BLOCK_SIZE); - PlatformVPipeBarrier(); -} - -// 业务子函数4：减法 + 指数 -// ls = ls - hm_block → ls = exp(ls) -__aicore__ __attribute__((always_inline)) inline void SubAndExp( - uint32_t sub_m, - uint32_t qk_n, - uint32_t qk_round_n -) -{ - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - - // *** ls = ls - hm_block - PlatformTensorSubValueRepeatM(ls32_ubuf_tensor, ls32_ubuf_tensor, - hm32_ubuf_tensor, tv32_ubuf_tensor, - sub_m, round_sub_m, qk_n, qk_round_n); - // *** ls = exp(ls) - PlatformExpV(ls32_ubuf_tensor, ls32_ubuf_tensor, - (sub_m * qk_round_n + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE); - PlatformVPipeBarrier(); -} - -// 业务子函数5：量化/转换输出 -// INT8: sub/exp/muls/brcb + QuantPerTokenImpl (FP32→INT8) -// 非INT8: conv_v (FP32→FP16) -// 共通: ub_to_gm(p_gm ← lp) + ReduceSum -__aicore__ __attribute__((always_inline)) inline void QuantizeAndOutput( - AscendC::GlobalTensor p_gm_tensor, - AscendC::LocalTensor ll_ubuf_tensor, - AscendC::LocalTensor pm32_ubuf_tensor, - uint32_t sub_m, - uint32_t qk_n, - uint32_t qk_round_n -) -{ - uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 128 - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - float quantMax = (float)1 / (float)127; - - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - PlatformSubV(pm32_ubuf_tensor, lm32_ubuf_tensor, hm32_ubuf_tensor, sub_m_d64); - PlatformVPipeBarrier(); - PlatformExpV(pm32_ubuf_tensor, pm32_ubuf_tensor, sub_m_d64); - PlatformVPipeBarrier(); - PlatformMulsV(pm32_ubuf_tensor, pm32_ubuf_tensor, quantMax, sub_m_d64); - PlatformVPipeBarrier(); - PlatformBrcbV(tv32_ubuf_tensor.template ReinterpretCast(), - pm32_ubuf_tensor.template ReinterpretCast(), round_sub_m / FLOAT_BLOCK_SIZE); - PlatformQuantPerToken(lp_ubuf_tensor, ls32_ubuf_tensor, tv32_ubuf_tensor, sub_m, qk_n, qk_round_n, 1); - } else { - PlatformConvVToOutput(lp_ubuf_tensor, ls32_ubuf_tensor, - (sub_m * qk_round_n + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE); - PlatformVPipeBarrier(); - } - PlatformVToMte3Sync(); - PlatformUbToGm(p_gm_tensor, lp_ubuf_tensor, - sub_m * qk_round_n * T_BLOCK_OFFSET / T_BLOCK_SIZE); - - // *** ll = rowsum(ls32) - PlatformReduceSumRepeatM(ll_ubuf_tensor, ls32_ubuf_tensor, sub_m, qk_n, qk_round_n); - PlatformVToMte2Notify(); - PlatformVPipeBarrier(); -} - -// 业务函数：非TP1 Phase 1 — SoftmaxStage1 编排（调用5个业务子函数） -__aicore__ __attribute__((always_inline)) inline void SoftmaxStage1( - AscendC::GlobalTensor p_gm_tensor, - AscendC::GlobalTensor s_gm_tensor, - AscendC::GlobalTensor s_rope_gm_tensor, - AscendC::GlobalTensor mask_gm_tensor, - AscendC::LocalTensor dm32_ubuf_tensor, - AscendC::LocalTensor ll_ubuf_tensor, - AscendC::LocalTensor pm32_ubuf_tensor, - uint32_t n_idx, - uint32_t qk_n, - uint32_t qk_round_n, - uint32_t sub_m, - uint32_t mask_offset, - const uint32_t sub_n_loop, - const uint32_t cur_batch, - const uint32_t start_kv, - const uint32_t real_n_loop, - const uint32_t head_idx, - const uint32_t pm_flag_scalar, - uint32_t cur_q_seqlen, - uint32_t cur_kv_seqlen, - bool need_mask -) -{ - // 段1：QK 数据加载（INT8:反量化+融合 / 非INT8:FP16+mask+Cast） - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - LoadQKDataInt8(s_gm_tensor, s_rope_gm_tensor, qk_n, qk_round_n, sub_m, head_idx); - } else { - LoadQKDataFP16(s_gm_tensor, mask_gm_tensor, qk_n, qk_round_n, sub_m, cur_q_seqlen, need_mask); - } - - // 段2：QK 缩放 + 非INT8 mask Add - ScaleAndMask(qk_n, qk_round_n, sub_m, cur_q_seqlen, need_mask); - - // 段3：Online Softmax 状态更新（ReduceMax + max/sub/ub_to_ub） - UpdateSoftmaxState(dm32_ubuf_tensor, n_idx, sub_m, qk_n, qk_round_n); - - // 段4：减法 + 指数（ls = ls - hm → ls = exp(ls)） - SubAndExp(sub_m, qk_n, qk_round_n); - - // 段5：量化/转换输出 + ub_to_gm + ReduceSum - QuantizeAndOutput(p_gm_tensor, ll_ubuf_tensor, pm32_ubuf_tensor, sub_m, qk_n, qk_round_n); -} - -// 业务函数：非TP1 Phase 1 — Softmax Stage1 调度 -// 包含 mask 计算、平台同步调用、SoftmaxStage1 调用（ping-pong） -__aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage1( - VectorContext &ctx, uint32_t n_idx, uint32_t start_kv) -{ - uint32_t cur_q_seqlen = ctx.cur_q_seqlen; - uint32_t cur_kv_seqlen = ctx.cur_kv_seqlen; - uint32_t cur_head_num = ctx.cur_head_num; - uint32_t cur_batch = ctx.cur_batch; - uint32_t pp_n_scalar = ctx.pp_n_scalar; - uint32_t sub_n_loop = ctx.sub_n_loop; - uint32_t real_n_loop = ctx.real_n_loop; - uint32_t n_loop = ctx.n_loop; - uint32_t sub_m = ctx.sub_m; - uint32_t head_idx = ctx.head_idx; - uint32_t mask_offset = ctx.mask_offset; - uint32_t tail_len = ctx.tail_len; - bool prev_tail_mask = ctx.prev_tail_mask; - - uint32_t qk_n = ctx.qk_n; - uint32_t qk_round_n = ctx.qk_round_n; - - // mask 边界判断 - bool need_mask = false; - uint32_t mask_start_offset = 0; - if (n_idx == (n_loop - 2)) { - need_mask = prev_tail_mask; - mask_start_offset = need_mask ? (tail_len + MASK_COLUMNS - 1) * MASK_COLUMNS : 0; - } - if (n_idx == (n_loop - 1)) { - qk_n = (cur_kv_seqlen - n_idx * pp_n_scalar); - qk_round_n = RoundUp<16>(qk_n); - need_mask = true; - mask_start_offset = (qk_n - 1) * MASK_COLUMNS; - } - - // 平台同步：SoftmaxStage1 前置同步（WaitFlagDev + WAIT_FLAG） - PlatformSoftmaxStage1PreSync(); - - // SoftmaxStage1 调用（ping-pong） - if (sub_m > 0) { - if (mask_type == 3) { - mask_start_offset = mask_offset + n_idx * pp_n_scalar; - } - // input QK shape (sub_m, qk_round_n) - if (n_idx % 2 == 0){ - SoftmaxStage1( - p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + - (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n * T_BLOCK_OFFSET + (uint64_t)(n_idx % 2) * TMP_SIZE * T_BLOCK_OFFSET / 2], - s_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + - (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - s_rope_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + - (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mask_gm_tensor[mask_start_offset], - dm32_ubuf_tensor, ll_ubuf_tensor, pm32_ubuf_tensor, - n_idx, qk_n, qk_round_n, sub_m, 0, sub_n_loop, cur_batch, start_kv, real_n_loop, head_idx, pm_flag_scalar1, cur_q_seqlen, cur_kv_seqlen, need_mask - ); - } else { - SoftmaxStage1( - p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + - (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n * T_BLOCK_OFFSET + - TMP_SIZE * T_BLOCK_OFFSET / 2], - s_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + - (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + - TMP_SIZE_DECODER / 2], - s_rope_gm_tensor[(int64_t)block_idx * TMP_SIZE_DECODER + - (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + - TMP_SIZE_DECODER / 2], - mask_gm_tensor[mask_start_offset], - dm32_stage2_ubuf_tensor, ll_stage2_ubuf_tensor, pm32_ubuf_stage2_tensor, - n_idx, qk_n, qk_round_n, sub_m, 0, sub_n_loop, cur_batch, start_kv, real_n_loop, head_idx, pm_flag_scalar2, cur_q_seqlen, cur_kv_seqlen, need_mask - ); - } - } - - // 平台同步：SoftmaxStage1 后置同步（FftsCrossCoreSync + SET_FLAG） - PlatformSoftmaxStage1PostSync(); -} - -// ==================== Stage2 业务子函数（三层拆分）==================== -// 以下子函数提取自 SoftmaxStage2MLAHeadLoop / TP1 / Tail 三个函数的共性段落， -// 通过平台原语下沉消除重复代码。 - -// 业务子函数：段落B — Merge Accumulate（brcb dm + load go + go=go*dm + go=lo+go） -// 非TP1/TP1/Tail 三函数完全一致的段落B逻辑 -__aicore__ __attribute__((always_inline)) inline void Stage2MergeAccumulate( - AscendC::GlobalTensor go_gm_tensor, - AscendC::LocalTensor go32_ubuf_tensor, - AscendC::LocalTensor lo_ubuf_tensor, - AscendC::LocalTensor tv32_ubuf_tensor, - AscendC::LocalTensor dm32_ubuf_tensor, - uint32_t sub_m, - uint32_t round_sub_m, - uint32_t round_v, - uint32_t head_loop) -{ - // brcb dm → tv, go = go * dm_block - PlatformMulVectorByBroadcast(go32_ubuf_tensor, tv32_ubuf_tensor.template ReinterpretCast(), - dm32_ubuf_tensor, sub_m, round_sub_m, round_v, __v); - - if (head_loop > 1) { - PlatformLoadGoFromGM(go_gm_tensor, go32_ubuf_tensor, sub_m, round_v); - } - - // go = lo + go - PlatformAddLoToGo(go32_ubuf_tensor, lo_ubuf_tensor, sub_m, round_v); -} - -// 业务子函数：段落D — Intermediate Save（保存 go32 到 GM，非最后一轮） -// 非TP1/TP1/Tail 三函数完全一致的段落D逻辑 -__aicore__ __attribute__((always_inline)) inline void Stage2IntermediateSave( - AscendC::GlobalTensor go_gm_tensor, - AscendC::LocalTensor go32_ubuf_tensor, - uint32_t sub_m, - uint32_t round_v) -{ - PlatformIntermediateSave(go_gm_tensor, go32_ubuf_tensor, sub_m, round_v); -} - -// 业务子函数：段落C — Finalize & Output（div + conv + DataCopyPad + Ring LSE） -// 非TP1/TP1 共用的最终输出逻辑（Tail 不使用此函数） -__aicore__ __attribute__((always_inline)) inline void Stage2FinalizeAndOutput( - AscendC::GlobalTensor o_gm_tensor, - AscendC::GlobalTensor go_gm_tensor, - AscendC::LocalTensor go32_ubuf_tensor, - AscendC::LocalTensor go_ubuf_tensor, - AscendC::LocalTensor tv32_ubuf_tensor, - AscendC::LocalTensor gl32_ubuf_tensor, - AscendC::LocalTensor lse32_ubuf_tensor, - AscendC::LocalTensor lse_conv_ubuf_tensor, - AscendC::LocalTensor gm32_ubuf_tensor, - AscendC::GlobalTensor lse_gm_tensor, - uint32_t sub_m, - uint32_t round_sub_m, - uint32_t round_v, - uint32_t head_loop, - uint32_t head_loop_idx, - uint32_t q_seq_len, - uint32_t numhead_per_process, - uint32_t head_res_row_num, - uint32_t head_start_sblock_idx, - uint32_t tail_res_row_num, - uint64_t o_offset) -{ - uint32_t sub_m_d64 = (sub_m + 63) / 64; - - // go = go / gl_block - PlatformDivVectorByBroadcast(go32_ubuf_tensor, tv32_ubuf_tensor.template ReinterpretCast(), - gl32_ubuf_tensor, sub_m, round_sub_m, round_v, __v, head_loop_idx); - - // go = castfp32to16(go) + V→MTE3 同步 - PlatformConvGoToOutput(go_ubuf_tensor, go32_ubuf_tensor, sub_m, round_v); - - // DataCopyPad output - PlatformDataCopyPadOutput(o_gm_tensor, go_ubuf_tensor, - sub_m, round_v, q_seq_len, numhead_per_process, - head_res_row_num, head_start_sblock_idx, tail_res_row_num); - - // Ring LSE copyout - if constexpr (IS_RING) { - PlatformRingLSECopyout(lse32_ubuf_tensor, gl32_ubuf_tensor, gm32_ubuf_tensor, - lse_conv_ubuf_tensor, lse_gm_tensor, sub_m, head_loop, o_offset); - } -} - -// 业务子函数：段落C — Tail Finalize & Output（仅 copyout gl/gm/go，无 div/conv） -// Tail 专用最终输出逻辑 -__aicore__ __attribute__((always_inline)) inline void TailStage2FinalizeAndOutput( - AscendC::GlobalTensor go_gm_tensor, - AscendC::GlobalTensor gl_gm_tensor, - AscendC::GlobalTensor gm_gm_tensor, - AscendC::LocalTensor go32_ubuf_tensor, - AscendC::LocalTensor gl32_ubuf_tensor, - AscendC::LocalTensor gm32_ubuf_tensor, - uint32_t sub_m, - uint32_t round_v) -{ - // Copyout gl/gm/go to GM - PlatformTailCopyout(gl_gm_tensor, gm_gm_tensor, go_gm_tensor, - gl32_ubuf_tensor, gm32_ubuf_tensor, go32_ubuf_tensor, - sub_m, round_v); -} - -// 业务函数：非TP1 Phase 2 — SoftmaxStage2MLAHeadLoop 核心 Head Loop 计算 -// 包含 load lo → DeQuant → exp(dm) → gl=dm*gl → gl=ll+gl → brcb dm → go=go*dm → go=lo+go -// → brcb gl → go=go/gl → conv_v → DataCopyPad输出 → Ring LSE copyout -__aicore__ __attribute__((always_inline)) inline void SoftmaxStage2MLAHeadLoop( - AscendC::GlobalTensor o_tmp_gm_tensor, - AscendC::GlobalTensor go_gm_tensor, - AscendC::GlobalTensor o_gm_tensor, - AscendC::LocalTensor dm32_ubuf_tensor, - AscendC::LocalTensor ll_ubuf_tensor, - AscendC::LocalTensor pm32_ubuf_tensor, - uint32_t n_idx, - uint32_t n_loop, - uint32_t qk_n, - uint32_t qk_round_n, - uint32_t sub_m, - uint64_t o_offset, - uint32_t head_idx, - uint32_t pm_flag_scalar, - uint32_t head_loop, - uint32_t head_loop_idx, - uint32_t q_seq_len, - uint32_t sub_head_num, - uint32_t cur_head_num, - uint32_t numhead_per_process, - uint32_t head_res_row_num, - uint32_t head_start_sblock_idx, - uint32_t tail_res_row_num - ) -{ - uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - PlatformVToMte2WaitEvent0(); - if (n_idx != 0) { - PlatformGmToUb( - lo_ubuf_tensor.template ReinterpretCast(), - o_tmp_gm_tensor, - sub_m * round_v / FLOAT_BLOCK_SIZE); - PlatformMte2ToVSync(); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - PlatformDeQuantPerHead( - deq_scale_gm_tensor_k1[head_idx], - o_tmp_gm_tensor, - lo_ubuf_tensor, lo_ubuf_tensor.template ReinterpretCast(), - descale_k1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, round_v, round_v, 1, 0); - } - } - PlatformSetVectorMaskFull(); - PlatformMte3ToMte2WaitEvent4(); - if (n_idx != 0) { - // *** dm = exp(dm) - if (head_loop_idx == 0) { - PlatformExpV(dm32_ubuf_tensor, dm32_ubuf_tensor, sub_m_d64); - PlatformVPipeBarrier(); - // *** gl = dm * gl - PlatformMulV(gl32_ubuf_tensor, dm32_ubuf_tensor, gl32_ubuf_tensor, sub_m_d64); - PlatformVPipeBarrier(); - // *** gl = ll + gl - PlatformAddV(gl32_ubuf_tensor, gl32_ubuf_tensor, ll_ubuf_tensor, sub_m_d64); - PlatformVPipeBarrier(); - } - // *** 段落B：brcb dm + load go + go=go*dm + go=lo+go - Stage2MergeAccumulate(go_gm_tensor, go32_ubuf_tensor, lo_ubuf_tensor, - tv32_ubuf_tensor, dm32_ubuf_tensor, - sub_m, round_sub_m, round_v, head_loop); - } else { - // *** gl = ll - if (head_loop_idx == 0) { - PlatformUbToUb(gl32_ubuf_tensor, ll_ubuf_tensor, 64 / FLOAT_BLOCK_SIZE); - PlatformVPipeBarrier(); - } - - PlatformGmToUb( - go32_ubuf_tensor.template ReinterpretCast(), - o_tmp_gm_tensor, - sub_m * round_v / FLOAT_BLOCK_SIZE); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - PlatformDeQuantPerHead( - deq_scale_gm_tensor_k1[head_idx], - o_tmp_gm_tensor, - go32_ubuf_tensor, go32_ubuf_tensor.template ReinterpretCast(), - descale_k1_ubuf_tensor, tv32_ubuf_tensor, pm32_ubuf_tensor, sub_m, round_v, round_v, 1, 0); - } else { - PlatformMte2ToVSync(); - } - } - PlatformVToMte2SetEvent0(); - - if (n_idx == n_loop - 1) { - // *** 段落C：brcb gl + go=go/gl + conv + DataCopyPad输出 + Ring LSE - Stage2FinalizeAndOutput(o_gm_tensor, go_gm_tensor, go32_ubuf_tensor, go_ubuf_tensor, - tv32_ubuf_tensor, gl32_ubuf_tensor, lse32_ubuf_tensor, lse_conv_ubuf_tensor, - gm32_ubuf_tensor, lse_gm_tensor, - sub_m, round_sub_m, round_v, head_loop, head_loop_idx, - q_seq_len, numhead_per_process, head_res_row_num, - head_start_sblock_idx, tail_res_row_num, o_offset); - } else if (head_loop > 1) { - // *** 段落D：中间结果保存 SET_FLAG + WAIT + ub_to_gm go32 - Stage2IntermediateSave(go_gm_tensor, go32_ubuf_tensor, sub_m, round_v); - } - PlatformMte3ToMte2SetEvent4(); -} - -// 业务函数：非TP1 Phase 2 — Softmax Stage2 调度 -// 包含 head 循环计算、WaitFlag、SoftmaxStage2MLAHeadLoop 调用 -__aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage2( - VectorContext &ctx, uint32_t n_idx) -{ - uint32_t cur_q_seqlen = ctx.cur_q_seqlen; - uint32_t cur_kv_seqlen = ctx.cur_kv_seqlen; - uint32_t cur_head_num = ctx.cur_head_num; - uint32_t pp_n_scalar = ctx.pp_n_scalar; - uint32_t n_loop = ctx.n_loop; - uint32_t sub_m = ctx.sub_m; - uint32_t sub_head_num = ctx.sub_head_num; - uint32_t head_idx = ctx.head_idx; - - uint32_t qk_n_2 = ctx.qk_n_2; - uint32_t qk_round_n_2 = ctx.qk_round_n_2; - - uint32_t process_row_num = 16; - uint32_t numhead_per_process = process_row_num / cur_q_seqlen; - - if (n_idx == n_loop) { - qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); - qk_round_n_2 = RoundUp(qk_n_2); - } - WaitFlagDev(UPDATE_READY_DECODER); - if (sub_m > 0) { - uint32_t head_loop = (sub_m + process_row_num - 1) / process_row_num; - - uint32_t head_res_row_num = 0; - uint32_t head_start_sblock_idx = 0; - uint32_t tail_res_row_num = 0; - - for (uint32_t head_loop_idx = 0; head_loop_idx < head_loop; ++head_loop_idx) { - uint32_t head_offset = head_loop_idx * process_row_num * round_v; - uint32_t cur_sub_m = head_loop_idx == (head_loop - 1) ? sub_m - head_loop_idx * process_row_num : process_row_num; // 15 or 3 - - // complete head num - head_start_sblock_idx = tail_res_row_num; - head_res_row_num = (cur_q_seqlen - tail_res_row_num) % cur_q_seqlen; - uint32_t cur_numhead_per_process = (cur_sub_m - head_res_row_num) / cur_q_seqlen; - tail_res_row_num = cur_sub_m - cur_numhead_per_process * cur_q_seqlen - head_res_row_num; - - uint32_t out_o_offset = head_loop_idx * numhead_per_process * round_v; // modified, round_v = 512 - - SoftmaxStage2MLAHeadLoop( - o_tmp_gm_tensor[(uint64_t)(block_idx * TMP_SIZE * 2 + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset + ((n_idx - 1) % 2) * TMP_SIZE)], - go_gm_tensor[(uint64_t)(block_idx * TMP_SIZE + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset)], - o_gm_tensor[(uint64_t)(o_offset + out_o_offset)], - dm32_ubuf_tensor[(uint64_t)((n_idx - 1) % 2 * 128 + head_loop_idx * process_row_num)], - ll_ubuf_tensor[(uint64_t)((n_idx - 1) % 2 * 256 + head_loop_idx * process_row_num)], - pm32_ubuf_tensor[(uint64_t)((n_idx - 1) % 2 * 128 + head_loop_idx * process_row_num)], - n_idx - 1, n_loop, qk_n_2, RoundUp(qk_round_n_2), cur_sub_m, o_offset, - head_idx + head_loop_idx * process_row_num, - pm_flag_scalar1, head_loop, head_loop_idx, cur_q_seqlen, sub_head_num, cur_head_num, - cur_numhead_per_process, - head_res_row_num, head_start_sblock_idx, tail_res_row_num); - } - } -} - -// 业务函数：非 TP1 — InnerRunVectorChange（内层业务编排） -// 初始化 VectorContext + n_loop 软流水循环（Stage1↔Stage2 交替） -__aicore__ __attribute__((always_inline)) inline void InnerRunVectorChange( - uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, - uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, - uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) -{ - VectorContext ctx; - InitVectorContext(ctx, cur_batch, start_head, cur_nIndx, - cur_q_seqlen, cur_kv_seqlen, cur_head_num, offset_tiling); - - o_offset = ctx.o_offset; - uint32_t start_kv = 0; - - for (uint32_t n_idx = 0; n_idx < ctx.n_loop + 1; n_idx++) { - if (n_idx != ctx.n_loop) { - ScheduleSoftmaxStage1(ctx, n_idx, start_kv); - } - if (n_idx != 0) { - ScheduleSoftmaxStage2(ctx, n_idx); - } - } -} - -// 业务函数：调度 Vector 非 TP1 任务（Run 方法的中段业务逻辑） -// 包含 batch/head 循环调度，读取 tiling 参数后调用 InnerRunVectorChange -__aicore__ __attribute__((always_inline)) inline void ScheduleVectorTasks() -{ - uint64_t cur_batch = 0; - uint32_t q_block_num_per_batch = (q_heads + cur_qn_blk_size - 1) / cur_qn_blk_size; - uint32_t process_num = q_block_num_per_batch * num_batches; - for (uint32_t process = block_idx; process < process_num; process += (uint32_t)block_num) { - cur_batch = process / q_block_num_per_batch; - if (cur_batch >= num_batches) break; - uint32_t offset_tiling = tiling_head_size + tiling_para_size * cur_batch; - uint32_t q_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); - if (kv_seqlen == 0) { - continue; - } - uint32_t start_head = (process % q_block_num_per_batch) * cur_qn_blk_size; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen; - uint32_t cur_head_num = cur_qn_blk_size; - uint32_t cur_nIndx = 0; - InnerRunVectorChange(cur_batch, start_head, cur_nIndx, cur_q_seq_len, cur_kv_seqlen, cur_head_num, - offset_tiling, 512, embed_split_loop_v_former); - } -} - -// ====== TP1 路径：Phase 1 / Phase 2 业务调度 ====== - -// 业务函数：TP1 Phase 1 — OnlineSoftmax Stage1 调度（InnerRunVectorChangeTP1 / TailInnerRunVectorChangeTP1 共用） -// 包含 tail qk_n 计算、WaitFlag、OnlineSoftmaxStage1 调用（m_ind 循环 ping-pong）、FftsCrossCoreSync -__aicore__ __attribute__((always_inline)) inline void ScheduleOnlineSoftmaxStage1( - VectorTP1Context &ctx, uint32_t n_idx) -{ - uint32_t cur_q_seqlen = ctx.cur_q_seqlen; - uint32_t cur_head_num = ctx.cur_head_num; - uint32_t cur_kv_seqlen = ctx.cur_kv_seqlen; - uint32_t pp_n_scalar = ctx.pp_n_scalar; - uint32_t sub_m = ctx.sub_m; - uint32_t n_loop = ctx.n_loop; - uint32_t s_block_stack = ctx.s_block_stack; - uint32_t m_slice = ctx.m_slice; - uint32_t m_end = ctx.m_end; - - uint32_t qk_n = ctx.qk_n; - uint32_t qk_round_n = ctx.qk_round_n; - - if (n_idx + s_block_stack > n_loop - 1) { - qk_n = (cur_kv_seqlen - n_idx * pp_n_scalar); - } else { - qk_n = pp_n_scalar * s_block_stack; - } - qk_round_n = RoundUp<16>(qk_n); - if (sub_m == 0) { - WaitFlagDev(QK_READY_DECODER); - } - uint32_t pingpong_flag = 0; - for (uint32_t m_ind = 0; m_ind < m_end; m_ind++) { - uint32_t row_offset = m_ind * m_slice; - uint32_t curr_m = m_ind == m_end - 1 ? sub_m - row_offset : m_slice; - uint32_t s_ub_offset = pingpong_flag * 8192; - uint32_t p_gm_offset = (uint64_t)block_idx * TMP_SIZE * 2 + - (uint64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + row_offset * qk_round_n + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE; - uint32_t s_gm_offset = (int64_t)block_idx * TMP_SIZE_DECODER * 4 + - (int64_t)sub_block_idx * cur_head_num * cur_q_seqlen / 2 * qk_round_n + row_offset * qk_round_n + (uint64_t)((n_idx / s_block_stack) % 2) * TMP_SIZE_DECODER * 2; - if (m_ind == 0) { - WaitFlagDev(QK_READY_DECODER); - } - if (curr_m == 0) { - continue; - } - OnlineSoftmaxStage1 ( - ls32_ubuf_tensor[s_ub_offset], - mask_ubuf_tensor, - mask_ubuf_tensor.template ReinterpretCast(), - lm32_ubuf_tensor[row_offset], - hm32_ubuf_tensor[row_offset], - gm32_ubuf_tensor[row_offset], - dm32_ubuf_tensor[((n_idx / s_block_stack) % 2) * UB_FLOAT_LINE_SIZE + row_offset], - ls32_ubuf_tensor[s_ub_offset], - ll_ubuf_tensor[row_offset], - gl32_ubuf_tensor[row_offset], - lp_ubuf_tensor[s_ub_offset * 2], - tv32_ubuf_tensor, - s_gm_tensor[s_gm_offset], - p_gm_tensor[p_gm_offset], - n_idx == 0, this->tor, - curr_m, qk_n, qk_round_n, pingpong_flag - ); - pingpong_flag = 1 - pingpong_flag; - } - FftsCrossCoreSync(SOFTMAX_READY_DECODER); -} - -// 业务函数：TP1 Phase 2 — SoftmaxStage2MLAHeadLoopTP1 核心 Head Loop 计算 -// 与非TP1相似，差异：n_idx!=4 vs n_idx!=0，无 DeQuant/exp(dm)/gl=dm*gl/gl=ll+gl 路径 -__aicore__ __attribute__((always_inline)) inline void SoftmaxStage2MLAHeadLoopTP1( - AscendC::GlobalTensor o_tmp_gm_tensor, - AscendC::GlobalTensor go_gm_tensor, - AscendC::GlobalTensor o_gm_tensor, - AscendC::LocalTensor dm32_ubuf_tensor, - AscendC::LocalTensor ll_ubuf_tensor, - AscendC::LocalTensor pm32_ubuf_tensor, - uint32_t n_idx, - uint32_t n_loop, - uint32_t qk_n, - uint32_t qk_round_n, - uint32_t sub_m, - uint64_t o_offset, - uint32_t head_idx, - uint32_t pm_flag_scalar, - uint32_t head_loop, - uint32_t head_loop_idx, - uint32_t q_seq_len, - uint32_t sub_head_num, - uint32_t cur_head_num, - uint32_t numhead_per_process, - uint32_t head_res_row_num, - uint32_t head_start_sblock_idx, - uint32_t tail_res_row_num - ) -{ - uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - PlatformVToMte2WaitEvent0(); - if (n_idx != 4) { - PlatformGmToUb( - lo_ubuf_tensor.template ReinterpretCast(), - o_tmp_gm_tensor, - sub_m * round_v / FLOAT_BLOCK_SIZE); - PlatformMte2ToVSync(); - } - PlatformSetVectorMaskFull(); - PlatformMte3ToMte2WaitEvent4(); - if (n_idx != 4) { - // expand_to_block - // *** 段落B：brcb dm + load go + go=go*dm + go=lo+go - Stage2MergeAccumulate(go_gm_tensor, go32_ubuf_tensor, lo_ubuf_tensor, - tv32_ubuf_tensor, dm32_ubuf_tensor, - sub_m, round_sub_m, round_v, head_loop); - } else { - // *** go = lo - - PlatformGmToUb( - go32_ubuf_tensor.template ReinterpretCast(), - o_tmp_gm_tensor, - sub_m * round_v / FLOAT_BLOCK_SIZE); - PlatformMte2ToVSync(); - } - PlatformVToMte2SetEvent0(); - - if (n_idx + 4 > n_loop + 4 - 1) { - // *** 段落C：brcb gl + go=go/gl + conv + DataCopyPad输出 + Ring LSE - Stage2FinalizeAndOutput(o_gm_tensor, go_gm_tensor, go32_ubuf_tensor, go_ubuf_tensor, - tv32_ubuf_tensor, gl32_ubuf_tensor, lse32_ubuf_tensor, lse_conv_ubuf_tensor, - gm32_ubuf_tensor, lse_gm_tensor, - sub_m, round_sub_m, round_v, head_loop, head_loop_idx, - q_seq_len, numhead_per_process, head_res_row_num, - head_start_sblock_idx, tail_res_row_num, o_offset); - } else if (head_loop > 1) { - // *** 段落D：中间结果保存 SET_FLAG + WAIT + ub_to_gm go32 - Stage2IntermediateSave(go_gm_tensor, go32_ubuf_tensor, sub_m, round_v); - } - PlatformMte3ToMte2SetEvent4(); -} - -// 业务函数：TP1 Phase 2 — SoftmaxStage2MLAHeadLoopTP1 调度（非 Tail 路径） -// 包含 head 循环计算、WaitFlag、SoftmaxStage2MLAHeadLoopTP1 调用 -__aicore__ __attribute__((always_inline)) inline void ScheduleSoftmaxStage2TP1( - VectorTP1Context &ctx, uint32_t n_idx) -{ - uint32_t cur_q_seqlen = ctx.cur_q_seqlen; - uint32_t cur_head_num = ctx.cur_head_num; - uint32_t cur_kv_seqlen = ctx.cur_kv_seqlen; - uint32_t pp_n_scalar = ctx.pp_n_scalar; - uint32_t sub_m = ctx.sub_m; - uint32_t sub_head_num = ctx.sub_head_num; - uint32_t n_loop = ctx.n_loop; - uint32_t head_idx = ctx.head_idx; - uint32_t s_block_stack = ctx.s_block_stack; - - uint32_t qk_n_2 = ctx.qk_n_2; - uint32_t qk_round_n_2 = ctx.qk_round_n_2; - - uint32_t process_row_num = 16; - uint32_t numhead_per_process = process_row_num / cur_q_seqlen; - - if (n_idx == n_loop) { - qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); - qk_round_n_2 = RoundUp(qk_n_2); - } - WaitFlagDev(UPDATE_READY_DECODER); - if (sub_m > 0) { - uint32_t head_loop = (sub_m + process_row_num - 1) / process_row_num; - - uint32_t head_res_row_num = 0; - uint32_t head_start_sblock_idx = 0; - uint32_t tail_res_row_num = 0; - - for (uint32_t head_loop_idx = 0; head_loop_idx < head_loop; ++head_loop_idx) { - uint32_t head_offset = head_loop_idx * process_row_num * round_v; - uint32_t cur_sub_m = head_loop_idx == (head_loop - 1) ? sub_m - head_loop_idx * process_row_num : process_row_num; - - // complete head num - head_start_sblock_idx = tail_res_row_num; - head_res_row_num = (cur_q_seqlen - tail_res_row_num) % cur_q_seqlen; - uint32_t cur_numhead_per_process = (cur_sub_m - head_res_row_num) / cur_q_seqlen; - tail_res_row_num = cur_sub_m - cur_numhead_per_process * cur_q_seqlen - head_res_row_num; - - uint32_t out_o_offset = head_loop_idx * numhead_per_process * round_v; - - SoftmaxStage2MLAHeadLoopTP1( - o_tmp_gm_tensor[(uint64_t)(block_idx * TMP_SIZE * 2 + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset + ((n_idx / s_block_stack - 1) % 2) * TMP_SIZE)], - go_gm_tensor[(uint64_t)(block_idx * TMP_SIZE + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset)], - o_gm_tensor[(uint64_t)(o_offset + out_o_offset)], - dm32_ubuf_tensor[(uint64_t)((n_idx / s_block_stack - 1) % 2 * UB_FLOAT_LINE_SIZE + head_loop_idx * process_row_num)], - ll_ubuf_tensor[(uint64_t)((n_idx / s_block_stack - 1) % 2 * 256 + head_loop_idx * process_row_num)], - pm32_ubuf_tensor, - n_idx, n_loop, qk_n_2, RoundUp(qk_round_n_2), cur_sub_m, o_offset, head_idx, - pm_flag_scalar1, head_loop, head_loop_idx, cur_q_seqlen, sub_head_num, cur_head_num, cur_numhead_per_process, - head_res_row_num, head_start_sblock_idx, tail_res_row_num); - } - } -} - -// 业务函数：TP1 Phase 2 (Tail) — TailSoftmaxStage2MLAHeadLoopTP1 核心 Head Loop 计算 -// 与 TP1 版本相似，但增加额外的 gl/gm/go copyout 逻辑 -__aicore__ __attribute__((always_inline)) inline void TailSoftmaxStage2MLAHeadLoopTP1( - AscendC::GlobalTensor o_tmp_gm_tensor, - AscendC::GlobalTensor go_gm_tensor, - AscendC::GlobalTensor gl_gm_tensor, - AscendC::GlobalTensor gm_gm_tensor, - AscendC::LocalTensor dm32_ubuf_tensor, - AscendC::LocalTensor go32_ubuf_tensor, - AscendC::LocalTensor gl32_ubuf_tensor, - AscendC::LocalTensor gm32_ubuf_tensor, - uint32_t n_idx, - uint32_t n_loop, - uint32_t qk_n, - uint32_t qk_round_n, - uint32_t sub_m, - uint64_t o_offset, - uint32_t head_idx, - uint32_t head_loop, - uint32_t head_loop_idx, - uint32_t q_seq_len, - uint32_t sub_head_num, - uint32_t cur_head_num, - uint32_t numhead_per_process - ) -{ - uint32_t sub_m_d64 = (sub_m + 63) / 64; // up aligned to 64 - uint32_t round_sub_m = (sub_m + 15) / 16 * 16; - PlatformVToMte2WaitEvent0(); - if (n_idx != 4) { - PlatformGmToUb( - lo_ubuf_tensor.template ReinterpretCast(), - o_tmp_gm_tensor, - sub_m * round_v / FLOAT_BLOCK_SIZE); - PlatformMte2ToVSync(); - } - PlatformSetVectorMaskFull(); - PlatformMte3ToMte2WaitEvent4(); - if (n_idx != 4) { - // expand_to_block - // *** 段落B：brcb dm + load go + go=go*dm + go=lo+go - Stage2MergeAccumulate(go_gm_tensor, go32_ubuf_tensor, lo_ubuf_tensor, - tv32_ubuf_tensor, dm32_ubuf_tensor, - sub_m, round_sub_m, round_v, head_loop); - } else { - // *** go = lo - - PlatformGmToUb( - go32_ubuf_tensor.template ReinterpretCast(), - o_tmp_gm_tensor, - sub_m * round_v / FLOAT_BLOCK_SIZE); - PlatformMte2ToVSync(); - } - PlatformVToMte2SetEvent0(); - - if (n_idx + 4 > n_loop + 4 - 1) { - // *** 段落C（Tail）：copyout gl/gm/go - TailStage2FinalizeAndOutput(go_gm_tensor, gl_gm_tensor, gm_gm_tensor, - go32_ubuf_tensor, gl32_ubuf_tensor, gm32_ubuf_tensor, - sub_m, round_v); - } else if (head_loop > 1) { - // *** 段落D：中间结果保存 SET_FLAG + WAIT + ub_to_gm go32 - Stage2IntermediateSave(go_gm_tensor, go32_ubuf_tensor, sub_m, round_v); - } - PlatformMte3ToMte2SetEvent4(); - PlatformPipeBarrierAll(); -} - -// 业务函数：TP1 Phase 2 (Tail) — TailSoftmaxStage2MLAHeadLoopTP1 调度（Tail 路径） -// 包含 head 循环计算、WaitFlag、TailSoftmaxStage2MLAHeadLoopTP1 调用 -__aicore__ __attribute__((always_inline)) inline void ScheduleTailSoftmaxStage2TP1( - VectorTP1Context &ctx, uint32_t n_idx) -{ - uint32_t cur_q_seqlen = ctx.cur_q_seqlen; - uint32_t cur_head_num = ctx.cur_head_num; - uint32_t cur_kv_seqlen = ctx.cur_kv_seqlen; - uint32_t pp_n_scalar = ctx.pp_n_scalar; - uint32_t sub_m = ctx.sub_m; - uint32_t sub_head_num = ctx.sub_head_num; - uint32_t n_loop = ctx.n_loop; - uint32_t head_idx = ctx.head_idx; - uint32_t start_head = ctx.start_head; - uint32_t s_block_stack = ctx.s_block_stack; - - uint32_t qk_n_2 = ctx.qk_n_2; - uint32_t qk_round_n_2 = ctx.qk_round_n_2; - - uint32_t process_row_num = 16; - uint32_t numhead_per_process = process_row_num / cur_q_seqlen; - - if (n_idx == n_loop) { - qk_n_2 = (cur_kv_seqlen - (n_idx - 1) * pp_n_scalar); - qk_round_n_2 = RoundUp(qk_n_2); - } - WaitFlagDev(UPDATE_READY_DECODER); - if (sub_m > 0) { - uint32_t head_loop = (sub_m + process_row_num - 1) / process_row_num; - - uint32_t head_res_row_num = 0; - uint32_t head_start_sblock_idx = 0; - uint32_t tail_res_row_num = 0; - - for (uint32_t head_loop_idx = 0; head_loop_idx < head_loop; ++head_loop_idx) { - uint32_t head_offset = head_loop_idx * process_row_num * round_v; - uint32_t cur_sub_m = head_loop_idx == (head_loop - 1) ? sub_m - head_loop_idx * process_row_num : process_row_num; - - // complete head num - head_start_sblock_idx = tail_res_row_num; - head_res_row_num = (cur_q_seqlen - tail_res_row_num) % cur_q_seqlen; - uint32_t cur_numhead_per_process = (cur_sub_m - head_res_row_num) / cur_q_seqlen; - tail_res_row_num = cur_sub_m - cur_numhead_per_process * cur_q_seqlen - head_res_row_num; - - uint32_t out_o_offset = head_loop_idx * numhead_per_process * round_v; - - TailSoftmaxStage2MLAHeadLoopTP1( - o_tmp_gm_tensor[(uint64_t)(block_idx * TMP_SIZE * 2 + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset + ((n_idx / s_block_stack - 1) % 2) * TMP_SIZE)], - go_gm_tensor[(uint64_t)(block_idx * TMP_SIZE + sub_block_idx * cur_head_num * cur_q_seqlen / 2 * round_v + head_offset)], - tmp_gm_tensor[(uint64_t)(block_idx * q_heads + start_head + sub_block_idx * cur_head_num * cur_q_seqlen / 2 + head_loop_idx * process_row_num)], - tmp_gm_tensor[(uint64_t)(block_num * q_heads + block_idx * q_heads + start_head + sub_block_idx * cur_head_num * cur_q_seqlen / 2 + head_loop_idx * process_row_num)], - dm32_ubuf_tensor[(uint64_t)((n_idx / s_block_stack - 1) % 2 * UB_FLOAT_LINE_SIZE + head_loop_idx * process_row_num)], - go32_ubuf_tensor, - gl32_ubuf_tensor[(uint64_t)(head_loop_idx * process_row_num)], - gm32_ubuf_tensor[(uint64_t)(head_loop_idx * process_row_num)], - n_idx, n_loop, qk_n_2, RoundUp(qk_round_n_2), cur_sub_m, o_offset, head_idx, - head_loop, head_loop_idx, cur_q_seqlen, sub_head_num, cur_head_num, cur_numhead_per_process - ); - } - } -} - -// 业务函数：TP1 — InnerRunVectorChangeTP1（内层业务编排） -// 初始化 VectorTP1Context + n_loop 软流水循环（Stage1↔Stage2TP1 交替） -__aicore__ __attribute__((always_inline)) inline void InnerRunVectorChangeTP1( - uint32_t cur_batch, uint32_t start_head, uint32_t cur_nIndx, - uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, - uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) -{ - VectorTP1Context ctx; - InitVectorTP1Context(ctx, cur_batch, start_head, cur_nIndx, - cur_q_seqlen, cur_kv_seqlen, cur_head_num, offset_tiling); - - o_offset = ctx.o_offset; - - for (uint32_t n_idx = 0; n_idx < ctx.n_loop + ctx.s_block_stack; n_idx += ctx.s_block_stack) { - if (n_idx < ctx.n_loop) { - ScheduleOnlineSoftmaxStage1(ctx, n_idx); - } - if (n_idx >= ctx.s_block_stack) { - ScheduleSoftmaxStage2TP1(ctx, n_idx); - } - } -} - -// 业务函数：TP1 Tail — TailInnerRunVectorChangeTP1（内层业务编排） -// 初始化 VectorTP1Context + n_loop 软流水循环（Stage1↔TailStage2TP1 交替） -__aicore__ __attribute__((always_inline)) inline void TailInnerRunVectorChangeTP1( - uint32_t start_head, - uint32_t cur_q_seqlen, uint32_t cur_kv_seqlen, uint32_t cur_head_num, - uint32_t offset_tiling, uint32_t embed_split_size_v, uint32_t embed_split_loop_v) -{ - VectorTP1Context ctx; - InitVectorTP1Context(ctx, 0, start_head, 0, - cur_q_seqlen, cur_kv_seqlen, cur_head_num, offset_tiling); - - o_offset = ctx.o_offset; - - for (uint32_t n_idx = 0; n_idx < ctx.n_loop + ctx.s_block_stack; n_idx += ctx.s_block_stack) { - if (n_idx < ctx.n_loop) { - ScheduleOnlineSoftmaxStage1(ctx, n_idx); - } - if (n_idx >= ctx.s_block_stack) { - ScheduleTailSoftmaxStage2TP1(ctx, n_idx); - } - } -} - -// ====== Tensor 逐行重复计算（Vector 业务基础函数）====== - -// Tensor 逐行减法（委托 arch 层平台函数） -__aicore__ __attribute__((always_inline)) inline void TensorSubValueRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& MaxTensor, - const AscendC::LocalTensor& tempMaxTensor, - uint32_t sub_m, - uint32_t round_sub_m, - uint32_t qk_n, - uint32_t qk_round_n) -{ - PlatformTensorSubValueRepeatM(dst, src, MaxTensor, tempMaxTensor, sub_m, round_sub_m, qk_n, qk_round_n); -} - -// Tensor 逐行除法（委托 arch 层平台函数） -__aicore__ __attribute__((always_inline)) inline void TensorDivRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& src1, - uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n) -{ - PlatformTensorDivRepeatM(dst, src, src1, sub_m, qk_n, qk_round_n); -} - -// Tensor 行归约求最大值（委托 arch 层平台函数） -__aicore__ __attribute__((always_inline)) inline void ReduceMaxRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& tempTensor, - uint32_t sub_m, - uint32_t qk_n, - uint32_t qk_round_n) -{ - PlatformReduceMaxRepeatM(dst, src, tempTensor, sub_m, qk_n, qk_round_n); -} - -// Tensor 行归约求和（委托 arch 层平台函数） -__aicore__ __attribute__((always_inline)) inline void ReduceSumRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - uint32_t sub_m, - uint32_t qk_n, - uint32_t qk_round_n) -{ - PlatformReduceSumRepeatM(dst, src, sub_m, qk_n, qk_round_n); -} - -// Tensor 逐行乘法（委托 arch 层平台函数） -__aicore__ __attribute__((always_inline)) inline void TensorMulRepeatM( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& src1, - uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t src1BlockStride -) { - PlatformTensorMulRepeatM(dst, src, src1, sub_m, qk_n, qk_round_n, src1BlockStride); -} - -// ====== 量化/反量化业务函数 ====== - -// ---- DeQuantPerHeadImpl 业务子函数 ---- - -// 业务子函数1：加载 deScale + online 乘 quantScale（委托arch层） -__aicore__ __attribute__((always_inline)) inline void LoadDeScaleAndOnlineMul( - const AscendC::GlobalTensor& deScaleGm, - AscendC::LocalTensor deScaleUb, - AscendC::LocalTensor quantScale, - uint32_t sub_m, bool online) -{ - PlatformLoadDeScaleAndOnlineMul(deScaleGm, deScaleUb, quantScale, sub_m, online); -} - -// 业务子函数2：加载 src(int32) + brcb 广播 deScale → tempScale（委托arch层） -__aicore__ __attribute__((always_inline)) inline void LoadSrcAndBrcbScale( - const AscendC::GlobalTensor& src, - AscendC::LocalTensor temp, - AscendC::LocalTensor deScaleUb, - AscendC::LocalTensor tempScale, - uint32_t sub_m, uint32_t qk_round_n, bool move_tensor) -{ - PlatformLoadSrcAndBrcbScale(src, temp, deScaleUb, tempScale, sub_m, qk_round_n, move_tensor); -} - -// 业务子函数3：INT32→FP32 转换(conv_v) + 乘 tempScale（委托arch层） -__aicore__ __attribute__((always_inline)) inline void ConvInt32ToFP32AndMul( - AscendC::LocalTensor dst, - AscendC::LocalTensor temp, - AscendC::LocalTensor tempScale, - uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n) -{ - PlatformConvInt32ToFP32AndMul(dst, temp, tempScale, sub_m, qk_n, qk_round_n); -} - -// 逐 Head 反量化：INT32 → FP32 + deScale 乘法（委托arch层编排） -__aicore__ __attribute__((always_inline)) inline void DeQuantPerHeadImpl( - const AscendC::GlobalTensor& deScaleGm, - const AscendC::GlobalTensor& src, - AscendC::LocalTensor dst, - AscendC::LocalTensor temp, - AscendC::LocalTensor deScaleUb, - AscendC::LocalTensor tempScale, - AscendC::LocalTensor quantScale, - uint32_t sub_m, - uint32_t qk_n, - uint32_t qk_round_n, - bool online, - bool move_tensor -){ - PlatformDeQuantPerHead(deScaleGm, src, dst, temp, deScaleUb, tempScale, quantScale, - sub_m, qk_n, qk_round_n, online, move_tensor); -} - -// 逐 Token 量化：FP32 → FP16 → INT8（委托arch层） -__aicore__ __attribute__((always_inline)) inline void QuantPerTokenImpl( - const AscendC::LocalTensor& dst, - const AscendC::LocalTensor& src, - const AscendC::LocalTensor& scale, - uint32_t sub_m, uint32_t qk_n, uint32_t qk_round_n, uint32_t pQuantOnline) -{ - PlatformQuantPerToken(dst, src, scale, sub_m, qk_n, qk_round_n, pQuantOnline); -} \ No newline at end of file diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_arch32.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_arch32.h deleted file mode 100644 index 19ebcc0..0000000 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_arch32.h +++ /dev/null @@ -1,640 +0,0 @@ -// Copyright 2025 The xLLM Authors. 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 -// -// https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE -// -// 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. -// ============================================================================== -// -// 文件：multi_latent_attention_arch32.h -// 层级：第三层 - 平台处理函数（Platform* 系列） -// 说明：本文件被 include 到 MLAttentionDecoderAic 类的 private 区域内， -// 包含所有与 Ascend C 平台 API 直接交互的函数（gm_to_l1 / l1_to_l0_a / -// l1_to_l0_b / l0c_to_gm / mmad / SET_FLAG / WAIT_FLAG / PIPE_BARRIER 等）。 -// 这些函数不包含业务逻辑，仅负责平台相关的数据搬运与同步操作。 - -// === 第三层：平台处理函数 === - -// 平台函数：设置 Q 数据搬运完成同步（MTE2→MTE1） -__aicore__ __attribute__((always_inline)) inline void PlatformSetQLoadComplete() -{ - SET_FLAG(MTE2, MTE1, EVENT_ID0); - WAIT_FLAG(MTE2, MTE1, EVENT_ID0); -} - -// 平台函数：等待 KV 数据搬运就绪（MTE1→MTE2，等待 V→L0B 完成） -__aicore__ __attribute__((always_inline)) inline void PlatformWaitKVLoadReady(uint32_t l1_kv_pingpong_flag) -{ - WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); -} - -// 平台函数：设置 KV 主体搬运完成 + 等待 V→L0B 完成（MTE2→MTE1 set, MTE1→MTE2 wait） -__aicore__ __attribute__((always_inline)) inline void PlatformSetKVMainLoadComplete(uint32_t l1_kv_pingpong_flag) -{ - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); -} - -// 平台函数：设置 KV Rope 搬运完成（MTE2→MTE1） -__aicore__ __attribute__((always_inline)) inline void PlatformSetKVRopeLoadComplete(uint32_t l1_kv_pingpong_flag) -{ - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); -} - -// 平台函数：Q 数据从 L1 加载到 L0A -__aicore__ __attribute__((always_inline)) inline void PlatformLoadQToL0A( - uint32_t embed_split_idx, uint32_t round_embed_split_size, - uint32_t q_load_coeff, bool is_rope) -{ - uint32_t loa_repeat = is_rope ? round_embed_split_size / BLOCK_SIZE - : round_embed_split_size / T_BLOCK_SIZE; - for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { - if (is_rope) { - l1_to_l0_a( - l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], - l1q_rope_buf_addr_tensor[loa_load_idx * CUBE_MATRIX_SIZE], - 0, loa_repeat, 0, q_load_coeff / BLOCK_SIZE, 0, 0); - } else { - l1_to_l0_a( - l0a_buf_tensor[embed_split_idx % 2 * 16384 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], - l1q_buf_addr_tensor[embed_split_idx * q_load_coeff * 128 + loa_load_idx * T_CUBE_MATRIX_SIZE], - 0, loa_repeat, 0, q_load_coeff / BLOCK_SIZE, 0, 0); - } - } -} - -// 平台函数：设置 Q 加载完成（MTE1→M） -__aicore__ __attribute__((always_inline)) inline void PlatformSetQLoadComplete( - uint32_t embed_split_idx) -{ - SET_FLAG(MTE1, M, embed_split_idx % 2); -} - -// 平台函数：等待 Q 加载完成（M→MTE1） -__aicore__ __attribute__((always_inline)) inline void PlatformWaitQLoadComplete( - uint32_t embed_split_idx) -{ - WAIT_FLAG(M, MTE1, embed_split_idx % 2); -} - -// 平台函数：KV 数据从 L1 加载到 L0B -__aicore__ __attribute__((always_inline)) inline void PlatformLoadKVToL0B( - uint32_t embed_split_idx, uint32_t round_embed_split_size, - uint32_t k_round_n, uint32_t qk_round_n, - uint32_t l1_kv_pingpong_flag, uint64_t hidden_size, bool is_rope) -{ - uint32_t l0b_repeat; - uint32_t l1_src_offset; - if (is_rope) { - l0b_repeat = round_embed_split_size * qk_round_n / CUBE_MATRIX_SIZE; - l1_src_offset = l1_kv_pingpong_flag * 128 * 64; - } else { - l0b_repeat = round_embed_split_size * k_round_n / T_CUBE_MATRIX_SIZE; - l1_src_offset = l1_kv_pingpong_flag * 128 * hidden_size + embed_split_idx * k_round_n * 128; - } - - if (embed_split_idx == 0) { - WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - } - if (embed_split_idx == 4) { - WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); - } - WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); - - if (is_rope) { - l1_to_l0_b( - l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], - l1kv_rope_buf_addr_tensor[l1_src_offset], - 0, l0b_repeat, 0, 1, 0, 0); - } else { - l1_to_l0_b( - l0b_buf_tensor[embed_split_idx % 2 * 16384], - l1kv_buf_addr_tensor[l1_src_offset], - 0, l0b_repeat, 0, 1, 0, 0); - } - - if (embed_split_idx == 4) { - SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); - } - SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); -} - -// 平台函数：QK MMAD 计算（含同步） -__aicore__ __attribute__((always_inline)) inline void PlatformComputeQKMMad( - uint32_t embed_split_idx, uint32_t embed_split_size, - uint32_t m_value, uint32_t qk_n, uint32_t qk_round_n_l1, - uint32_t l1_kv_pingpong_flag, bool is_rope) -{ - WAIT_FLAG(MTE1, M, embed_split_idx % 2); - WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); - if (embed_split_idx == 0) { - WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); - } - - if (is_rope) { - mmad( - mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], - l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], - l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], - m_value, qk_n, embed_split_size, 1); - } else { - uint32_t n_value = qk_n; - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - n_value = qk_round_n_l1; - } - mmad( - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - l0a_buf_tensor[embed_split_idx % 2 * 16384], - l0b_buf_tensor[embed_split_idx % 2 * 16384], - m_value, n_value, embed_split_size, embed_split_idx == 0); - } - - PIPE_BARRIER(M); - SET_FLAG(M, MTE1, embed_split_idx % 2); - SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); -} - -// 平台函数：L0C 结果拷贝到 GM（含同步） -__aicore__ __attribute__((always_inline)) inline void PlatformCopyQKResultToGM( - uint32_t embed_split_idx, uint32_t m_value, uint32_t qk_n, - uint32_t qk_round_n, uint32_t l1_kv_pingpong_flag, - uint32_t n_idx, bool is_rope) -{ - if (is_rope) { - SET_FLAG(M, FIX, l1_kv_pingpong_flag); - WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - l0c_to_gm( - s_rope_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], - m_value, qk_round_n, RoundUp<16>(m_value), qk_round_n); - } else { - l0c_to_gm( - s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - m_value, qk_round_n, RoundUp<16>(m_value), qk_round_n); - } - SET_FLAG(FIX, M, l1_kv_pingpong_flag); - return; - } - - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - if (embed_split_idx == 3) { - SET_FLAG(M, FIX, l1_kv_pingpong_flag); - WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); - l0c_to_gm( - s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - m_value, qk_n, RoundUp<16>(m_value), qk_round_n); - SET_FLAG(FIX, M, l1_kv_pingpong_flag); - } - } - if (embed_split_idx == 4) { - SET_FLAG(M, FIX, l1_kv_pingpong_flag); - WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); - l0c_to_gm( - s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - m_value, qk_round_n, RoundUp<16>(m_value), qk_round_n); - SET_FLAG(FIX, M, l1_kv_pingpong_flag); - } -} - -// 平台函数：KV L1→L0B 转置加载（k_round_n <= round_embed_split_size 分支） -__aicore__ __attribute__((always_inline)) inline void PlatformLoadKVTransposeSmallN( - uint32_t l0b_pingpong_flag, uint32_t l1_kv_pingpong_flag, - uint64_t l1kv_offset, uint64_t k_round_n, uint64_t hidden_size, - uint32_t embed_split_size, uint32_t round_embed_split_size) -{ - WAIT_FLAG(M, MTE1, l0b_pingpong_flag + 2); - AscendC::LoadData2dTransposeParams loadDataParams; - loadDataParams.dstGap = 0; - loadDataParams.startIndex = 0; - loadDataParams.dstFracGap = 0; - loadDataParams.repeatTimes = round_embed_split_size / T_BLOCK_SIZE; - loadDataParams.srcStride = k_round_n / T_BLOCK_SIZE; - uint16_t dstGap = sizeof(IN_DTYPE) == 1 ? 1 : 0; - loadDataParams.dstGap = dstGap; - for (uint32_t l0b_load_idx = 0; l0b_load_idx < k_round_n / T_BLOCK_SIZE; ++l0b_load_idx) { - AscendC::LoadDataWithTranspose( - l0b_buf_tensor[l0b_pingpong_flag * 16384 + l0b_load_idx * RoundUp<16>(embed_split_size) * T_BLOCK_SIZE], - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + l1kv_offset + l0b_load_idx * T_BLOCK_SIZE * T_BLOCK_SIZE], - loadDataParams); - } -} - -// 平台函数：KV L1→L0B 转置加载（k_round_n > round_embed_split_size 分支） -__aicore__ __attribute__((always_inline)) inline void PlatformLoadKVTransposeLargeN( - uint32_t l0b_pingpong_flag, uint32_t l1_kv_pingpong_flag, - uint64_t l1kv_offset, uint64_t hidden_size, - uint32_t qk_round_n_2, uint32_t round_embed_split_size) -{ - WAIT_FLAG(M, MTE1, l0b_pingpong_flag + 2); - AscendC::LoadData2dTransposeParams loadDataParams; - loadDataParams.dstGap = 0; - loadDataParams.startIndex = 0; - loadDataParams.dstFracGap = 0; - for (uint32_t l0b_load_idx = 0; l0b_load_idx < round_embed_split_size / T_BLOCK_SIZE; ++l0b_load_idx) { - loadDataParams.repeatTimes = qk_round_n_2 / T_BLOCK_SIZE; - loadDataParams.srcStride = 1; - loadDataParams.dstGap = round_embed_split_size / BLOCK_SIZE - 1; - AscendC::LoadDataWithTranspose( - l0b_buf_tensor[l0b_pingpong_flag * 16384 + l0b_load_idx * T_BLOCK_SIZE * T_BLOCK_SIZE], - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * hidden_size + l1kv_offset + l0b_load_idx * qk_round_n_2 * T_BLOCK_SIZE], - loadDataParams); - } -} - -// 平台函数：KV 加载完成后设置 MTE1→MTE2 同步 -__aicore__ __attribute__((always_inline)) inline void PlatformSetKVLoadComplete( - bool is_last_split, uint32_t l1_kv_pingpong_flag) -{ - if (is_last_split) { - SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); - } -} - -// 平台函数：P 数据从 GM 加载到 L1 -__aicore__ __attribute__((always_inline)) inline void PlatformLoadPFromGMToL1( - uint32_t n_idx, uint32_t row_num, uint64_t k_round_n, - uint32_t qk_round_n_2) -{ - WaitFlagDev(SOFTMAX_READY_DECODER); - WAIT_FLAG(MTE1, MTE2, EVENT_ID7); - gm_to_l1( - l1p_buf_addr_tensor, - p_gm_tensor[(uint64_t)block_idx * TMP_SIZE * T_BLOCK_OFFSET + ((n_idx - 1) % 2) * TMP_SIZE * T_BLOCK_OFFSET / 2], - row_num, - RoundUp(row_num), - 0, - k_round_n, - 0, - qk_round_n_2 * 2 / sizeof(IN_DTYPE)); - SET_FLAG(MTE2, MTE1, EVENT_ID7); - WAIT_FLAG(MTE2, MTE1, EVENT_ID7); -} - -// 平台函数：P 数据从 L1 加载到 L0A（INT8 路径） -__aicore__ __attribute__((always_inline)) inline void PlatformLoadPToL0AInt8( - uint32_t l0_p_pingpong_flag, uint32_t row_num, - uint32_t qk_round_n_2_l1) -{ - WAIT_FLAG(M, MTE1, l0_p_pingpong_flag); - l1_to_l0_a( - l0a_buf_tensor[l0_p_pingpong_flag * 16384], l1p_buf_addr_tensor, - RoundUp(row_num), qk_round_n_2_l1, 0, 0, 0, 0); -} - -// 平台函数：P 数据从 L1 加载到 L0A（非 INT8 路径） -__aicore__ __attribute__((always_inline)) inline void PlatformLoadPToL0AGeneral( - uint32_t l0_p_pingpong_flag, uint32_t row_num, - uint32_t qk_round_n_2) -{ - WAIT_FLAG(M, MTE1, l0_p_pingpong_flag); - uint32_t p_load_coeff = RoundUp<16>(row_num); - for (uint64_t loa_load_idx = 0; loa_load_idx < p_load_coeff / BLOCK_SIZE; ++loa_load_idx) { - l1_to_l0_a( - l0a_buf_tensor[l0_p_pingpong_flag * 16384 + loa_load_idx * qk_round_n_2 * BLOCK_SIZE], - l1p_buf_addr_tensor[loa_load_idx * T_CUBE_MATRIX_SIZE], - 0, qk_round_n_2 / T_BLOCK_SIZE, 0, p_load_coeff / BLOCK_SIZE, 0, 0); - } -} - -// 平台函数：P 加载完成后设置 MTE1→MTE2 同步 -__aicore__ __attribute__((always_inline)) inline void PlatformSetPLoadComplete() -{ - SET_FLAG(MTE1, MTE2, EVENT_ID7); -} - -// 平台函数：PV MMAD 计算（含同步） -__aicore__ __attribute__((always_inline)) inline void PlatformComputePVMmad( - uint32_t l0b_pingpong_flag, uint32_t l0c_pingpong_flag, - uint32_t l0_p_pingpong_flag, uint32_t embed_split_size, - uint32_t qk_n_2, bool is_last_split) -{ - SET_FLAG(MTE1, M, l0b_pingpong_flag); - WAIT_FLAG(MTE1, M, l0b_pingpong_flag); - WAIT_FLAG(FIX, M, l0c_pingpong_flag); - mmad( - mm2_l0c_buf_tensor[l0c_pingpong_flag * 16384], - l0a_buf_tensor[l0_p_pingpong_flag * 16384], - l0b_buf_tensor[l0b_pingpong_flag * 16384], - m, embed_split_size, qk_n_2, 1); - SET_FLAG(M, MTE1, l0b_pingpong_flag + 2); - if (is_last_split) { - SET_FLAG(M, MTE1, l0_p_pingpong_flag); - } -} - -// 平台函数：L0C 结果拷贝到 GM（含同步） -__aicore__ __attribute__((always_inline)) inline void PlatformCopyPVResultToGM( - uint32_t l0c_pingpong_flag, uint32_t embed_split_idx, - uint32_t n_idx, uint32_t round_embed_split_size) -{ - SET_FLAG(M, FIX, l0c_pingpong_flag); - WAIT_FLAG(M, FIX, l0c_pingpong_flag); - l0c_to_gm( - o_tmp_gm_tensor[(uint64_t)block_idx * TMP_SIZE * 2 + embed_split_idx * round_embed_split_size + ((n_idx - 1) % 2) * TMP_SIZE], - mm2_l0c_buf_tensor[l0c_pingpong_flag * 16384], - m, RoundUp<16>(round_embed_split_size), RoundUp<16>(m), round_v); - SET_FLAG(FIX, M, l0c_pingpong_flag); -} - -// ==================== Rope 平台函数（INT8 专有）==================== - -// 平台函数：Q_Rope 从 L1 加载到 L0A（for 循环 l1_to_l0_a） -__aicore__ __attribute__((always_inline)) inline void PlatformLoadQRopeToL0A( - uint32_t embed_split_idx, uint32_t q_load_coeff, - uint32_t round_embed_split_size) -{ - WAIT_FLAG(M, MTE1, embed_split_idx % 2); - for (uint64_t loa_load_idx = 0; loa_load_idx < q_load_coeff / BLOCK_SIZE; ++loa_load_idx) { - l1_to_l0_a( - l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2 + loa_load_idx * round_embed_split_size * BLOCK_SIZE], - l1q_rope_buf_addr_tensor[loa_load_idx * CUBE_MATRIX_SIZE], - 0, - round_embed_split_size / BLOCK_SIZE, - 0, - q_load_coeff / BLOCK_SIZE, - 0, - 0 - ); - } - SET_FLAG(MTE1, M, embed_split_idx % 2); -} - -// 平台函数：KV_Rope 从 L1 加载到 L0B（单次 l1_to_l0_b + 同步） -__aicore__ __attribute__((always_inline)) inline void PlatformLoadKVRopeToL0B( - uint32_t embed_split_idx, uint32_t l1_kv_pingpong_flag, - uint32_t round_embed_split_size, uint32_t qk_round_n) -{ - WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag + 2); - WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); - l1_to_l0_b( - l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], - l1kv_rope_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64], - 0, - round_embed_split_size * qk_round_n / CUBE_MATRIX_SIZE, - 0, - 1, - 0, - 0 - ); - SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag + 2); - SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); -} - -// 平台函数：Rope MMAD 计算（含同步） -__aicore__ __attribute__((always_inline)) inline void PlatformComputeQRopeMMad( - uint32_t embed_split_idx, uint32_t l1_kv_pingpong_flag, - uint32_t m_value, uint32_t qk_n, uint32_t embed_split_size) -{ - WAIT_FLAG(MTE1, M, embed_split_idx % 2); - WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); - } - mmad( - mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], - l0a_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], - l0b_buf_tensor.template ReinterpretCast()[embed_split_idx % 2 * 16384 * 2], - m_value, - qk_n, - embed_split_size, - 1 - ); - PIPE_BARRIER(M); - SET_FLAG(M, MTE1, embed_split_idx % 2); - SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); -} - -// 平台函数：Rope 结果 L0C 拷贝到 GM（含同步） -__aicore__ __attribute__((always_inline)) inline void PlatformCopyQRopeResultToGM( - uint32_t l1_kv_pingpong_flag, uint32_t n_idx, - uint32_t m_value, uint32_t qk_round_n) -{ - SET_FLAG(M, FIX, l1_kv_pingpong_flag); - WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - l0c_to_gm( - s_rope_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor.template ReinterpretCast()[l1_kv_pingpong_flag * 16384], - m_value, - qk_round_n, - RoundUp<16>(m_value), - qk_round_n - ); - } else { - l0c_to_gm( - s_gm_tensor[(uint64_t)block_idx * TMP_SIZE_DECODER + (uint64_t)(n_idx % 2) * TMP_SIZE_DECODER / 2], - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - m_value, - qk_round_n, - RoundUp<16>(m_value), - qk_round_n - ); - } - SET_FLAG(FIX, M, l1_kv_pingpong_flag); -} - -// 平台函数：初始化管道同步事件（Run 方法头部批量 SET_FLAG） -__aicore__ __attribute__((always_inline)) inline void PlatformInitPipeSync() -{ - SET_FLAG(M, MTE1, EVENT_ID0); - SET_FLAG(M, MTE1, EVENT_ID1); - SET_FLAG(M, MTE1, EVENT_ID2); - SET_FLAG(M, MTE1, EVENT_ID3); - SET_FLAG(M, MTE1, EVENT_ID4); - SET_FLAG(M, MTE1, EVENT_ID5); - SET_FLAG(M, MTE1, EVENT_ID6); - SET_FLAG(M, MTE1, EVENT_ID7); - SET_FLAG(FIX, M, EVENT_ID0); - SET_FLAG(FIX, M, EVENT_ID1); - SET_FLAG(MTE1, MTE2, EVENT_ID0); - SET_FLAG(MTE1, MTE2, EVENT_ID1); - SET_FLAG(MTE1, MTE2, EVENT_ID2); - SET_FLAG(MTE1, MTE2, EVENT_ID3); - SET_FLAG(MTE1, MTE2, EVENT_ID4); - SET_FLAG(MTE1, MTE2, EVENT_ID5); - SET_FLAG(MTE1, MTE2, EVENT_ID6); - SET_FLAG(MTE1, MTE2, EVENT_ID7); - SET_FLAG(FIX, MTE1, EVENT_ID0); - SET_FLAG(FIX, MTE1, EVENT_ID1); - SET_FLAG(FIX, MTE1, EVENT_ID2); - SET_FLAG(FIX, MTE1, EVENT_ID3); - SET_FLAG(FIX, MTE1, EVENT_ID4); - SET_FLAG(FIX, MTE1, EVENT_ID5); - SET_FLAG(MTE2, FIX, EVENT_ID0); -} - -// 平台函数：等待管道同步完成（Run 方法尾部批量 WAIT_FLAG + BARRIER） -__aicore__ __attribute__((always_inline)) inline void PlatformWaitPipeSync() -{ - WAIT_FLAG(M, MTE1, EVENT_ID0); - WAIT_FLAG(M, MTE1, EVENT_ID1); - WAIT_FLAG(M, MTE1, EVENT_ID2); - WAIT_FLAG(M, MTE1, EVENT_ID3); - WAIT_FLAG(M, MTE1, EVENT_ID4); - WAIT_FLAG(M, MTE1, EVENT_ID5); - WAIT_FLAG(M, MTE1, EVENT_ID6); - WAIT_FLAG(M, MTE1, EVENT_ID7); - WAIT_FLAG(FIX, M, EVENT_ID0); - WAIT_FLAG(FIX, M, EVENT_ID1); - WAIT_FLAG(MTE1, MTE2, EVENT_ID0); - WAIT_FLAG(MTE1, MTE2, EVENT_ID1); - WAIT_FLAG(MTE1, MTE2, EVENT_ID2); - WAIT_FLAG(MTE1, MTE2, EVENT_ID3); - WAIT_FLAG(MTE1, MTE2, EVENT_ID4); - WAIT_FLAG(MTE1, MTE2, EVENT_ID5); - WAIT_FLAG(MTE1, MTE2, EVENT_ID6); - WAIT_FLAG(MTE1, MTE2, EVENT_ID7); - WAIT_FLAG(FIX, MTE1, EVENT_ID0); - WAIT_FLAG(FIX, MTE1, EVENT_ID1); - WAIT_FLAG(FIX, MTE1, EVENT_ID2); - WAIT_FLAG(FIX, MTE1, EVENT_ID3); - WAIT_FLAG(FIX, MTE1, EVENT_ID4); - WAIT_FLAG(FIX, MTE1, EVENT_ID5); - WAIT_FLAG(MTE2, FIX, EVENT_ID0); - PIPE_BARRIER(ALL); -} - -// ==================== TP1 QK 平台函数 ==================== - -// 平台函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2 分支） -// TP1 专有：L1 目标偏移 [l1_kv_pingpong_flag * 128 * 256]，dValue=256，使用 block_table_id+embed_split_idx 寻址 -__aicore__ __attribute__((always_inline)) inline void PlatformLoadTP1KVMainToL1( - uint32_t embed_split_idx, uint32_t qk_n, uint32_t qk_round_n, - uint32_t l1_kv_pingpong_flag, int64_t kv_offset, uint32_t embed_split_idx_128) -{ - WAIT_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); - if constexpr (KInputType == InputFormat::ND_FORMAT) { - gm_to_l1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 256], - k_gm_tensor[kv_offset + embed_split_idx_128], - qk_n, - qk_round_n, - 0, - 256, - 0, - stride_kv); - } else { - gm_to_l1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 256], - k_gm_tensor[kv_offset + block_size * 128 * embed_split_idx], - qk_round_n, - block_size, - qk_round_n, - 256, - 256, 256); - } - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); -} - -// 平台函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4 分支） -// TP1 专有：L1 目标偏移 [l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128]，dValue=64 -__aicore__ __attribute__((always_inline)) inline void PlatformLoadTP1KVRopeToL1( - uint32_t qk_n, uint32_t qk_round_n, - uint32_t l1_kv_pingpong_flag, int64_t kv_offset_rope) -{ - WAIT_FLAG(MTE1, MTE2, 2 + l1_kv_pingpong_flag); - if constexpr (KInputType == InputFormat::ND_FORMAT) { - gm_to_l1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128], - k_rope_gm_tensor[kv_offset_rope], - qk_n, - qk_round_n, - 0, - 64, - 0, - stride_kv_rope); - } else { - gm_to_l1( - l1kv_buf_addr_tensor[l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128], - k_rope_gm_tensor[kv_offset_rope], - qk_round_n, - block_size, - qk_round_n, - 64, - 64, - 64); - } - SET_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); - WAIT_FLAG(MTE2, MTE1, l1_kv_pingpong_flag); -} - -// 平台函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） -// TP1 专有：L1 源偏移 [now_l1_offset + embed_split_idx % 2 * qk_round_n * 128] -__aicore__ __attribute__((always_inline)) inline void PlatformLoadTP1KVToL0B( - uint32_t embed_split_idx, uint32_t round_embed_split_size, - uint32_t qk_round_n, int64_t now_l1_offset, uint32_t l1_kv_pingpong_flag) -{ - WAIT_FLAG(M, MTE1, embed_split_idx % 2 + 2); - l1_to_l0_b( - l0b_buf_tensor[embed_split_idx % 2 * 16384], - l1kv_buf_addr_tensor[now_l1_offset + embed_split_idx % 2 * qk_round_n * 128], - 0, - round_embed_split_size * qk_round_n / T_CUBE_MATRIX_SIZE, - 0, - 1, - 0, - 0); - // 释放 KV L1 缓冲区 - if (embed_split_idx == 1 || embed_split_idx == 3) { - SET_FLAG(MTE1, MTE2, l1_kv_pingpong_flag); - } - if (embed_split_idx == 4) { - SET_FLAG(MTE1, MTE2, 2 + l1_kv_pingpong_flag); - } - SET_FLAG(MTE1, M, embed_split_idx % 2 + 2); -} - -// 平台函数：TP1 QK MMAD 计算（含同步） -// TP1 与非 TP1 的 cmatrixInitVal 均为 (embed_split_idx == 0)，可复用逻辑 -__aicore__ __attribute__((always_inline)) inline void PlatformComputeTP1QKMMad( - uint32_t embed_split_idx, uint32_t embed_split_size, - uint32_t m_value, uint32_t qk_n, uint32_t l1_kv_pingpong_flag) -{ - WAIT_FLAG(MTE1, M, embed_split_idx % 2); - WAIT_FLAG(MTE1, M, embed_split_idx % 2 + 2); - if (embed_split_idx == 0) { - WAIT_FLAG(FIX, M, l1_kv_pingpong_flag); - } - mmad( - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - l0a_buf_tensor[embed_split_idx % 2 * 16384], - l0b_buf_tensor[embed_split_idx % 2 * 16384], - m_value, - qk_n, - embed_split_size, - embed_split_idx == 0); - PIPE_BARRIER(M); - SET_FLAG(M, MTE1, embed_split_idx % 2); - SET_FLAG(M, MTE1, embed_split_idx % 2 + 2); -} - -// 平台函数：TP1 QK 结果从 L0C 拷贝到 GM（embed_split_idx == 4 时） -// TP1 专有：GM 目标偏移 [block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar] -// dstStride = sv_round_n -__aicore__ __attribute__((always_inline)) inline void PlatformCopyTP1QKResultToGM( - uint32_t m_value, uint32_t qk_round_n, uint32_t l1_kv_pingpong_flag, - uint64_t gm_dst_offset, uint32_t sv_round_n) -{ - SET_FLAG(M, FIX, l1_kv_pingpong_flag); - WAIT_FLAG(M, FIX, l1_kv_pingpong_flag); - l0c_to_gm( - s_gm_tensor[gm_dst_offset], - mm1_l0c_buf_tensor[l1_kv_pingpong_flag * 16384], - m_value, - qk_round_n, - RoundUp<16>(m_value), - sv_round_n); - SET_FLAG(FIX, M, l1_kv_pingpong_flag); -} \ No newline at end of file diff --git a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h b/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h deleted file mode 100644 index 7ca42de..0000000 --- a/xllm_ops/multi_latent_attention/op_kernel/multi_latent_attention_bs.h +++ /dev/null @@ -1,696 +0,0 @@ -// Copyright 2025 The xLLM Authors. 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 -// -// https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE -// -// 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. -// ============================================================================== -// -// 文件：multi_latent_attention_bs.h -// 层级：第二层 - 业务逻辑函数（Load*/Init*/Compute*/Copy* 系列） -// 说明：本文件被 include 到 MLAttentionDecoderAic 类的 private 区域内， -// 包含所有纯流程控制函数，这些函数负责调用第三层 Platform* 系列函数 -// 完成具体的业务编排，不直接调用平台 API。 -// 依赖：第三层函数（Platform* 系列）来自 multi_latent_attention_arch32.h - -// === 第二层：业务逻辑函数 === - -// 业务函数：将 Q 主体数据从 GM 搬运到 L1（ND→NZ 格式转换） -// 覆盖三种场景： -// 1. cur_q_seqlen == 1 → gm_to_l1 单矩阵搬运 -// 2. cur_q_seqlen > 1 && q_heads < 128 → DataCopy 批量多矩阵搬运 -// 3. cur_q_seqlen > 1 && q_heads >= 128 → for循环逐token搬运（规避 stride 位宽限制） -__aicore__ __attribute__((always_inline)) inline void LoadQMainFromGMToL1( - AscendC::LocalTensor &l1_dst, - AscendC::GlobalTensor &gm_src, - uint64_t src_offset, - uint32_t cur_q_seqlen, - uint32_t cur_head_num) -{ - if (cur_q_seqlen == 1) { - gm_to_l1( - l1_dst, - gm_src[src_offset], - cur_head_num, // nValue - RoundUp<16>(cur_head_num), // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 512, // dValue - 0, // dstNzMatrixStride, unused - 512 // srcDValue - ); - } else { - if (q_heads < 128) { - AscendC::DataCopy( - l1_dst, - gm_src[src_offset], - AscendC::Nd2NzParams( - cur_q_seqlen, // ndNum - cur_head_num, // nValue - 512, // dValue - 512 * q_heads, // srcNdMatrixStride - 512, // srcDValue - RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride - cur_q_seqlen, // dstNzNStride - 16 // dstNzMatrixStride - ) - ); - } else { - for (uint32_t ii = 0; ii < cur_q_seqlen; ii++) { - AscendC::DataCopy( - l1_dst[ii * 16], // offset one datablock - gm_src[src_offset + ii * q_heads * 512], - AscendC::Nd2NzParams( - 1, // ndNum - cur_head_num, // nValue - 512, // dValue - 0, // srcNdMatrixStride - 512, // srcDValue - RoundUp<16>(cur_q_seqlen * cur_head_num), // dstNzC0Stride - cur_q_seqlen, // dstNzNStride - 16 // dstNzMatrixStride - ) - ); - } - } - } -} - -// ==================== TP1 QK 业务函数 ==================== - -// TP1 QK 参数结构体 -struct TP1QKParams { - uint32_t qk_n; - uint32_t qk_round_n; - uint32_t l1_kv_pingpong_flag; - uint32_t embed_split_size; - uint32_t round_embed_split_size; - int64_t now_l1_offset; - int64_t kv_offset; - int64_t kv_offset_rope; - uint32_t sv_round_n; - uint32_t gm_split_idx; // split_idx 用于 GM 偏移计算 - uint32_t pp_n_scalar; // block_size，用于 GM 偏移计算 -}; - -// 业务函数：初始化 TP1 QK 参数 -// 对应原始 CUBE1 stage1 中 split_idx 循环体内的参数计算 -__aicore__ __attribute__((always_inline)) inline void InitTP1QKParams( - TP1Context &ctx, uint32_t split_idx, uint32_t n_idx, TP1QKParams ¶ms) -{ - uint32_t pp_n_scalar = ctx.pp_n_scalar; - uint32_t n_loop = ctx.n_loop; - uint32_t now_idx = n_idx + split_idx; - - params.l1_kv_pingpong_flag = now_idx % 2; - - // 动态计算 qk_n / qk_round_n - if (now_idx == (n_loop - 1)) { - params.qk_n = (ctx.cur_kv_seqlen - now_idx * pp_n_scalar); - } else { - params.qk_n = pp_n_scalar; - } - params.qk_round_n = RoundUp(params.qk_n); - - // embed_split_size 默认 128，idx==4 时为 64（在编排层设置） - params.embed_split_size = 128; - params.round_embed_split_size = RoundUp(params.embed_split_size); - - // now_l1_offset 默认 0，在 LoadTP1KVMain/KVRope 中更新 - params.now_l1_offset = 0; - - // block_table_id / kv_offset / kv_offset_rope 地址计算 - uint32_t block_table_id = (uint32_t)(*(block_tables_gm + - ctx.cur_batch * max_num_blocks_per_query + ctx.start_kv / block_size + now_idx)); - params.kv_offset = (int64_t)block_table_id * block_size * stride_kv; - params.kv_offset_rope = (int64_t)block_table_id * block_size * stride_kv_rope; - - // sv_round_n 计算（用于 S→GM 的 dstStride） - uint32_t sv_n_triu = n_loop * pp_n_scalar; - uint32_t sv_n; - if (n_idx + ctx.s_block_stack > n_loop - 1) { - sv_n = ctx.cur_kv_seqlen - n_idx * pp_n_scalar; - } else { - sv_n = pp_n_scalar * ctx.s_block_stack; - } - params.sv_round_n = (sv_n + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE; - - params.gm_split_idx = split_idx; - params.pp_n_scalar = pp_n_scalar; -} - -// 业务函数：TP1 Q 数据从 L1 加载到 L0A -// 复用非 TP1 的 PlatformWaitQLoadComplete / PlatformLoadQToL0A / PlatformSetQLoadComplete -// TP1 中 q_load_coeff = m，与非 TP1 一致 -__aicore__ __attribute__((always_inline)) inline void LoadTP1QDataToL0A( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformWaitQLoadComplete(embed_split_idx); - PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, m, false); - PlatformSetQLoadComplete(embed_split_idx); -} - -// 业务函数：TP1 KV Main 数据从 GM 搬运到 L1（idx 0,2） -// 更新 now_l1_offset 供后续 L1→L0B 使用 -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVMainFromGM( - TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 256; - PlatformLoadTP1KVMainToL1(embed_split_idx, params.qk_n, params.qk_round_n, - params.l1_kv_pingpong_flag, params.kv_offset, - embed_split_idx * 128); -} - -// 业务函数：TP1 KV Rope 数据从 GM 搬运到 L1（idx 4） -// 更新 now_l1_offset 供后续 L1→L0B 使用 -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVRopeFromGM( - TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - params.now_l1_offset = params.l1_kv_pingpong_flag * 128 * 64 + 2 * 256 * 128; - PlatformLoadTP1KVRopeToL1(params.qk_n, params.qk_round_n, - params.l1_kv_pingpong_flag, params.kv_offset_rope); -} - -// 业务函数：TP1 KV 数据从 L1 加载到 L0B（所有 embed_split_idx） -__aicore__ __attribute__((always_inline)) inline void LoadTP1KVDataToL0B( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformLoadTP1KVToL0B(embed_split_idx, params.round_embed_split_size, - params.qk_round_n, params.now_l1_offset, - params.l1_kv_pingpong_flag); -} - -// 业务函数：TP1 QK MMAD 计算 -__aicore__ __attribute__((always_inline)) inline void ComputeTP1QKMMad( - const TP1QKParams ¶ms, uint32_t embed_split_idx) -{ - PlatformComputeTP1QKMMad(embed_split_idx, params.embed_split_size, - m, params.qk_n, params.l1_kv_pingpong_flag); -} - -// 业务函数：TP1 QK 结果拷贝到 GM（idx 4 时） -// GM 偏移: block_idx * TMP_SIZE_DECODER * 4 + ((n_idx/s_block_stack)%2) * TMP_SIZE_DECODER * 2 + split_idx * pp_n_scalar -__aicore__ __attribute__((always_inline)) inline void CopyTP1QKResultToGM( - const TP1QKParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) -{ - uint64_t gm_dst_offset = (uint64_t)block_idx * TMP_SIZE_DECODER * 4 + - (uint64_t)((n_idx / 4) % 2) * TMP_SIZE_DECODER * 2 + - params.gm_split_idx * params.pp_n_scalar; - PlatformCopyTP1QKResultToGM(m, params.qk_round_n, - params.l1_kv_pingpong_flag, - gm_dst_offset, params.sv_round_n); -} - -// 业务函数：将 Q Rope 数据从 GM 搬运到 L1（ND→NZ 格式转换） -// INT8 场景：用 gm_to_l1 搬到独立的 l1q_rope_buf_addr_tensor -// 非INT8 场景：用 DataCopy 搬到 l1q_buf_addr_tensor 的 Q 主体之后 -__aicore__ __attribute__((always_inline)) inline void LoadQRopeFromGMToL1( - AscendC::LocalTensor &l1_q, - AscendC::LocalTensor &l1_q_rope, - AscendC::GlobalTensor &gm_src, - uint64_t src_offset, - uint32_t cur_q_seqlen, - uint32_t cur_head_num) -{ - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - gm_to_l1( - l1_q_rope, - gm_src[src_offset], - cur_head_num, // nValue - RoundUp<16>(cur_head_num), // dstNzC0Stride - 0, // dstNzMatrixStride, unused - 64, // dValue - 0, // dstNzMatrixStride, unused - 64 // srcDValue - ); - } else { - AscendC::DataCopy( - l1_q[RoundUp<16>(cur_head_num * cur_q_seqlen) * 512], - gm_src[src_offset], - AscendC::Nd2NzParams( - cur_head_num, // ndNum - cur_q_seqlen, // nValue - 64, // dValue - 64, // srcNdMatrixStride - 64 * q_heads, // srcDValue - RoundUp<16>(cur_head_num * cur_q_seqlen), // dstNzC0Stride - 1, // dstNzNStride - 16 * cur_q_seqlen // dstNzMatrixStride - ) - ); - } -} - - -// 业务函数：将 KV 主体数据从 GM 搬运到 L1 -// ND_FORMAT 场景：ND→NZ 格式转换（gm_to_l1） -// INT8 / 非INT8 NZ 场景：NZ→NZ 格式搬运（gm_to_l1） -template -__aicore__ __attribute__((always_inline)) inline void LoadKVMainFromGMToL1( - AscendC::LocalTensor l1_dst, - AscendC::GlobalTensor gm_src, - uint32_t n_value, // 实际行数 - uint32_t dst_nz_c0_stride, // L1 对齐行数（dstNzC0Stride） - uint32_t d_value, // 列数（512） - uint32_t src_d_value, // GM 行 stride（ND场景为stride_kv，NZ场景为0） - bool is_nd_to_nz) // true=ND→NZ, false=NZ→NZ -{ - if (is_nd_to_nz) { - gm_to_l1( - l1_dst, - gm_src, - n_value, // nValue - dst_nz_c0_stride, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - d_value, // dValue - 0, // dstNzMatrixStride, unused - src_d_value // srcDValue - ); - } else { - gm_to_l1( - l1_dst, - gm_src, - n_value, // nValue - dst_nz_c0_stride, // dstNzC0Stride - 0, // dstNzMatrixStride, unused - d_value, // dValue - 0, // dstNzMatrixStride, unused - 0 // srcDValue (NZ→NZ always 0) - ); - } -} - - -// 业务函数：将 KV Rope 数据从 GM 搬运到 L1 -// ND_FORMAT 场景：ND→NZ 格式转换（gm_to_l1），d_value=64, src_d_value=stride_kv_rope -// INT8 / 非INT8 NZ 场景：NZ→NZ 格式搬运（gm_to_l1），d_value=64, src_d_value=0 -template -__aicore__ __attribute__((always_inline)) inline void LoadKVRopeFromGMToL1( - AscendC::LocalTensor l1_dst, - AscendC::GlobalTensor gm_src, - uint32_t n_value, // 实际行数 - uint32_t dst_nz_c0_stride, // L1 对齐行数（dstNzC0Stride） - uint32_t d_value, // 列数（64） - uint32_t src_d_value, // GM 行 stride（ND场景为stride_kv_rope，NZ场景为0） - bool is_nd_to_nz) // true=ND→NZ, false=NZ→NZ -{ - // 逻辑与 LoadKVMainFromGMToL1 完全一致，仅参数不同（d_value=64） - LoadKVMainFromGMToL1(l1_dst, gm_src, n_value, dst_nz_c0_stride, d_value, src_d_value, is_nd_to_nz); -} - -// 业务函数：初始化 QK 参数 -__aicore__ __attribute__((always_inline)) inline void InitQKParams( - MLAContext &ctx, uint32_t l1_kv_pingpong_flag, QKParams ¶ms) -{ - params.qk_n = ctx.qk_n; - params.qk_round_n = ctx.qk_round_n; - params.qk_round_n_l1 = ctx.qk_round_n_l1; - params.hidden_size = ctx.hidden_size; - params.k_round_n = ctx.k_round_n; - params.row_num = ctx.row_num; - params.embed_split_size = 128; - params.round_embed_split_size = RoundUp(params.embed_split_size); - params.q_load_coeff = m; - params.hidden_split_time = (params.hidden_size + 128 - 1) / 128; - params.l1_kv_pingpong_flag = l1_kv_pingpong_flag; -} - -// 业务函数：Q 数据从 L1 加载到 L0A -__aicore__ __attribute__((always_inline)) inline void LoadQDataToL0A( - const QKParams ¶ms, uint32_t embed_split_idx, bool is_rope) -{ - PlatformWaitQLoadComplete(embed_split_idx); - PlatformLoadQToL0A(embed_split_idx, params.round_embed_split_size, - params.q_load_coeff, is_rope); - PlatformSetQLoadComplete(embed_split_idx); -} - -// 业务函数：KV 数据从 L1 加载到 L0B -__aicore__ __attribute__((always_inline)) inline void LoadKVDataToL0B( - const QKParams ¶ms, uint32_t embed_split_idx, bool is_rope) -{ - PlatformLoadKVToL0B(embed_split_idx, params.round_embed_split_size, - params.k_round_n, params.qk_round_n, - params.l1_kv_pingpong_flag, params.hidden_size, is_rope); -} - -// 业务函数：QK MMAD 计算 -__aicore__ __attribute__((always_inline)) inline void ComputeQKMMad( - const QKParams ¶ms, uint32_t embed_split_idx, bool is_rope) -{ - PlatformComputeQKMMad(embed_split_idx, params.embed_split_size, - m, params.qk_n, params.qk_round_n_l1, - params.l1_kv_pingpong_flag, is_rope); -} - -// 业务函数：QK 结果拷贝到 GM -__aicore__ __attribute__((always_inline)) inline void CopyQKResultToGM( - const QKParams ¶ms, uint32_t embed_split_idx, - uint32_t n_idx, bool is_rope) -{ - PlatformCopyQKResultToGM(embed_split_idx, m, params.qk_n, - params.qk_round_n, params.l1_kv_pingpong_flag, - n_idx, is_rope); -} - -// 业务函数：初始化 PV 参数 -__aicore__ __attribute__((always_inline)) inline void InitPVParams( - MLAContext &ctx, uint32_t n_idx, PVParams ¶ms) -{ - params.qk_n_2 = ctx.qk_n_2; - params.qk_round_n_2 = ctx.qk_round_n_2; - params.qk_round_n_2_l1 = ctx.qk_round_n_2_l1; - params.k_round_n = ctx.k_round_n; - params.row_num = ctx.row_num; - params.hidden_size = ctx.hidden_size; - - if (n_idx == ctx.n_loop) { - params.qk_n_2 = (ctx.cur_kv_seqlen - (n_idx - 1) * ctx.pp_n_scalar); - params.qk_round_n_2 = RoundUp(params.qk_n_2); - params.qk_round_n_2_l1 = RoundUp(params.qk_n_2); - } - params.k_round_n = params.qk_round_n_2_l1; - params.l1_kv_pingpong_flag = (n_idx - 1) % 2; - params.l0_p_pingpong_flag = (n_idx - 1) % 2; - params.embed_split_size = 128; - embed_split_loop_v = 4; - params.round_embed_split_size = RoundUp(params.embed_split_size); -} - -// 业务函数：KV L1→L0B 转置加载 -__aicore__ __attribute__((always_inline)) inline void LoadKVTransposeToL0B( - const PVParams ¶ms, uint32_t embed_split_idx) -{ - bool is_last_split = (embed_split_idx == embed_split_loop_v - 1); - if (params.k_round_n <= params.round_embed_split_size) { - PlatformLoadKVTransposeSmallN( - params.l0b_pingpong_flag, params.l1_kv_pingpong_flag, - params.l1kv_offset, params.k_round_n, params.hidden_size, - params.embed_split_size, params.round_embed_split_size); - } else { - PlatformLoadKVTransposeLargeN( - params.l0b_pingpong_flag, params.l1_kv_pingpong_flag, - params.l1kv_offset, params.hidden_size, - params.qk_round_n_2, params.round_embed_split_size); - } - PlatformSetKVLoadComplete(is_last_split, params.l1_kv_pingpong_flag); -} - -// 业务函数：P 数据加载（GM→L1→L0A） -__aicore__ __attribute__((always_inline)) inline void LoadPDataToL0A( - const PVParams ¶ms, uint32_t n_idx) -{ - PlatformLoadPFromGMToL1( - n_idx, params.row_num, params.k_round_n, params.qk_round_n_2); - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - PlatformLoadPToL0AInt8( - params.l0_p_pingpong_flag, params.row_num, - params.qk_round_n_2_l1); - } else { - PlatformLoadPToL0AGeneral( - params.l0_p_pingpong_flag, params.row_num, - params.qk_round_n_2); - } - PlatformSetPLoadComplete(); -} - -// 业务函数：PV MMAD 计算 -__aicore__ __attribute__((always_inline)) inline void ComputePVMmad( - const PVParams ¶ms, uint32_t embed_split_idx) -{ - bool is_last_split = (embed_split_idx == embed_split_loop_v - 1); - PlatformComputePVMmad( - params.l0b_pingpong_flag, params.l0c_pingpong_flag, - params.l0_p_pingpong_flag, params.embed_split_size, - params.qk_n_2, is_last_split); -} - -// 业务函数：PV 结果拷贝到 GM -__aicore__ __attribute__((always_inline)) inline void CopyPVResultToGM( - const PVParams ¶ms, uint32_t embed_split_idx, uint32_t n_idx) -{ - PlatformCopyPVResultToGM( - params.l0c_pingpong_flag, embed_split_idx, - n_idx, params.round_embed_split_size); -} - -// ==================== Rope 业务函数（INT8 专有）==================== - -// 业务函数：Rope 计算（Q_Rope × KV_Rope） -__aicore__ __attribute__((always_inline)) inline void ComputeQRope( - const QKParams ¶ms, uint32_t n_idx) -{ - if constexpr (tilingKeyType == TilingKeyType::TILING_INT8_DATA) { - // Rope 部分固定使用 embed_split_idx=4, embed_split_size=64 - uint32_t rope_embed_split_idx = 4; - uint32_t rope_embed_split_size = 64; - uint32_t rope_round_embed_split_size = 64; - - // 1. 加载 Q_Rope 到 L0A - PlatformLoadQRopeToL0A( - rope_embed_split_idx, params.q_load_coeff, - rope_round_embed_split_size); - - // 2. 加载 KV_Rope 到 L0B - PlatformLoadKVRopeToL0B( - rope_embed_split_idx, params.l1_kv_pingpong_flag, - rope_round_embed_split_size, params.qk_round_n); - - // 3. mmad 矩阵乘 - PlatformComputeQRopeMMad( - rope_embed_split_idx, params.l1_kv_pingpong_flag, - m, params.qk_n, rope_embed_split_size); - - // 4. 结果搬回 GM - PlatformCopyQRopeResultToGM( - params.l1_kv_pingpong_flag, n_idx, - m, params.qk_round_n); - } -} - - -// 业务函数：调度 Cube MLA 任务（Run 方法的 for 循环部分） -// 遍历所有 batch，解析 tiling 参数，调用 InnerRunCubeMLA -__aicore__ __attribute__((always_inline)) inline void ScheduleCubeTasks() -{ - uint32_t q_block_num_per_batch = (q_heads + cur_qn_blk_size - 1) / cur_qn_blk_size; - uint32_t process_num = q_block_num_per_batch * num_batches; - - for (uint32_t process = block_idx; process < process_num; process += (uint32_t)block_num) { - uint32_t cur_batch = process / q_block_num_per_batch; - if (cur_batch >= num_batches) break; - - uint32_t offset_tiling = tiling_head_size + tiling_para_size * cur_batch; - uint32_t start_core_idx = (cur_batch * q_block_num_per_batch) % block_num; - - uint32_t q_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 1 + offset_tiling)); - if (kv_seqlen == 0) { - continue; - } - uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; - - uint32_t start_head = (process % q_block_num_per_batch) * cur_qn_blk_size; - uint32_t start_kv = 0; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen; - uint32_t cur_head_num = cur_qn_blk_size; - - InnerRunCubeMLA(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, - offset_tiling); - } -} - - -// 业务函数：调度 Cube MLA TP1 任务（RunTP1 方法的中段业务逻辑） -// 包含主循环调度 + tail 优化三分支（cores_per_seq 动态调整） -__aicore__ __attribute__((always_inline)) inline void ScheduleCubeTasksTP1() -{ - uint32_t tail = totalTaskNum % block_num; - if constexpr (EnableOptimization) { - - } else{ - tail = 0; // control whether to run tail optimization - } - uint32_t totalTaskNumRound = totalTaskNum - tail; - - - for (uint32_t process = block_idx; process < totalTaskNumRound; process += (uint32_t)block_num) { // for task - uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; - uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - - uint32_t q_seqlen = 1; - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); - if (kv_seqlen == 0) { - continue; - } - uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; - - uint32_t start_head = 0; - uint32_t start_kv = 0; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen; - uint32_t cur_head_num = q_heads; - - InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); - } - - // suppose all seqs have same length - if (tail > 0){ - uint32_t sample_kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + tiling_head_size + 2)); - bool enableExtraOptimization = true; - if (block_num % 4 == 3) { - // cannot optimize this situation due to math problem - enableExtraOptimization = false; - } - if (sample_kv_seqlen <= 2048){ - // Too Short to benefit from optimization - enableExtraOptimization = false; - } - if (!enableExtraOptimization || tail <= block_num / 2) { - // collect all metadata - uint32_t cores_per_seq = 1; - if (0 < tail && tail <= block_num / 4) {// 6 tasks left, each works with 4 cores - cores_per_seq = 4; - if (tail == 1){ - cores_per_seq = block_num; - } - else if (tail == 2){ - cores_per_seq = block_num / 2; - } - else if(tail == 3){ - cores_per_seq = block_num / 3; - } - else if(tail == 4){ - cores_per_seq = block_num / 4; - } - } - else if(block_num / 4 < tail && tail <= block_num / 3) { // 8 tasks left, each works with 3 cores - cores_per_seq = 3; - - } - else if(block_num / 3 < tail && tail <= block_num / 2) { // 12 tasks left, each works with 2 cores - cores_per_seq = 2; - } - else { - // no extra optimization for tail > 12 - cores_per_seq = 1; - } - - if(!enableExtraOptimization){ - cores_per_seq = 1; - } - - uint32_t process = totalTaskNumRound + block_idx / cores_per_seq; - if (process < totalTaskNum) { - uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; - uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - - uint32_t q_seqlen = 1; - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); - uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; - uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; - uint32_t actual_work_cores = kv_seqlen / kv_seqlen_align + (kv_seqlen % kv_seqlen_align != 0); - // cores_per_seq = actual_work_cores; - uint32_t kv_seqlen_process = 0; - if (block_idx < block_idx / cores_per_seq * cores_per_seq + actual_work_cores){ - kv_seqlen_process = (block_idx % cores_per_seq == actual_work_cores - 1) ? - (kv_seqlen - kv_seqlen_align * (actual_work_cores - 1)) : kv_seqlen_align; - } - - if (kv_seqlen > 0 && kv_seqlen_process > 0) { - uint32_t start_head = 0; - uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen_process; - uint32_t cur_head_num = q_heads; - - // no need to modify anything in cube kernel, just call the same kernel - InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); - } - } - } - else if (tail > 3 * block_num / 4){ - // no benefit for optimizing this situation - uint32_t process = totalTaskNumRound + block_idx; - if (process < totalTaskNum) { - uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; - uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - - uint32_t q_seqlen = 1; - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); - if (kv_seqlen > 0) { - uint32_t kv_seqlen_align = (kv_seqlen + block_size - 1) / block_size * block_size; - - uint32_t start_head = 0; - uint32_t start_kv = 0; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen; - uint32_t cur_head_num = q_heads; - - InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); - } - } - } - else { - // 18 >= tail >= 12 - // first 12 tasks, two cores per task - { - uint32_t cores_per_seq = 2; - uint32_t process = totalTaskNumRound + block_idx / cores_per_seq; - uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; - uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - - uint32_t q_seqlen = 1; - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); - uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; - uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; - uint32_t kv_seqlen_process = (block_idx % cores_per_seq == cores_per_seq - 1) ? - (kv_seqlen - kv_seqlen_align * (cores_per_seq - 1)) : kv_seqlen_align; - - if (kv_seqlen > 0 && kv_seqlen_process > 0) { - uint32_t start_head = 0; - uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen_process; - uint32_t cur_head_num = q_heads; - - // no need to modify anything in cube kernel, just call the same kernel - InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); - } - } - { - uint32_t cores_per_seq = 4; - uint32_t process = totalTaskNumRound + block_num / 2 + block_idx / cores_per_seq; - if (process < totalTaskNum) { - uint32_t offset_tiling = tiling_head_size + tiling_para_size * process; - uint32_t cur_batch = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + offset_tiling)); - - uint32_t q_seqlen = 1; - uint32_t kv_seqlen = (uint32_t)(*((__gm__ uint32_t *)tiling_gm + 2 + offset_tiling)); - uint32_t kv_seqlen_each = kv_seqlen / cores_per_seq; - uint32_t kv_seqlen_align = (kv_seqlen_each + block_size - 1) / block_size * block_size; - uint32_t kv_seqlen_process = (block_idx % cores_per_seq == cores_per_seq - 1) ? - (kv_seqlen - kv_seqlen_align * (cores_per_seq - 1)) : kv_seqlen_align; - - if (kv_seqlen > 0 && kv_seqlen_process > 0) { - uint32_t start_head = 0; - uint32_t start_kv = (block_idx % cores_per_seq) * kv_seqlen_align; - uint32_t cur_q_seq_len = q_seqlen; - uint32_t cur_kv_seqlen = kv_seqlen_process; - uint32_t cur_head_num = q_heads; - - // no need to modify anything in cube kernel, just call the same kernel - InnerRunCubeMLATP1(cur_batch, start_head, cur_head_num, start_kv, cur_q_seq_len, cur_kv_seqlen, offset_tiling); - } - } - } - } - } -} From da9e7ff2112a072d5648effd60634e9aa6d0af46 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Sat, 15 Aug 2026 05:55:14 +0800 Subject: [PATCH 20/24] feat: slove the exp bug --- .../multi_latent_attention/op_kernel/mixkernels/include/simd.h | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/simd.h b/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/simd.h index 3947bd9..6df1d44 100644 --- a/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/simd.h +++ b/xllm_ops/multi_latent_attention/op_kernel/mixkernels/include/simd.h @@ -189,7 +189,8 @@ __aicore__ inline void exp_v(AscendC::LocalTensor dst, #if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) // 3510: Use PRECISION_1ULP_FTZ_FALSE to preserve Subnormal numbers, // improving online softmax precision for large kv_seqlen. - AscendC::Exp( + static constexpr AscendC::ExpConfig MLA_EXP_CFG{AscendC::ExpAlgo::PRECISION_1ULP_FTZ_FALSE}; + AscendC::Exp( dst, src, (uint64_t)0, From ad4f45bbce5aafdd6f87aedcb4cecf018cf59a3d Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Tue, 18 Aug 2026 15:19:44 +0800 Subject: [PATCH 21/24] feat: x_attention adapter the a5 --- xllm_ops/build_aclnn.sh | 2 +- xllm_ops/x_attention/op_host/CMakeLists.txt | 35 ++ .../op_host/x_attention_tiling.cpp | 261 +++++++++- .../x_attention/op_host/x_attention_tiling.h | 67 +++ xllm_ops/x_attention/op_host/xa_arch_config.h | 21 + .../x_attention/op_host/xa_arch_config.h.in | 21 + .../op_kernel/arch35/combine_kernel.h | 150 ++++++ .../arch35/shared_infer_catlass_kernel.h | 444 ++++++++++++++++++ .../arch35/unshared_infer_catlass_kernel.h | 375 +++++++++++++++ .../arch35/x_attention_catlass_helper.h | 120 +++++ .../op_kernel/arch35/x_attention_common.h | 109 +++++ .../x_attention/op_kernel/x_attention.cpp | 82 +++- 12 files changed, 1658 insertions(+), 29 deletions(-) create mode 100644 xllm_ops/x_attention/op_host/xa_arch_config.h create mode 100644 xllm_ops/x_attention/op_host/xa_arch_config.h.in create mode 100644 xllm_ops/x_attention/op_kernel/arch35/combine_kernel.h create mode 100644 xllm_ops/x_attention/op_kernel/arch35/shared_infer_catlass_kernel.h create mode 100644 xllm_ops/x_attention/op_kernel/arch35/unshared_infer_catlass_kernel.h create mode 100644 xllm_ops/x_attention/op_kernel/arch35/x_attention_catlass_helper.h create mode 100644 xllm_ops/x_attention/op_kernel/arch35/x_attention_common.h diff --git a/xllm_ops/build_aclnn.sh b/xllm_ops/build_aclnn.sh index e62f2fc..a0136d7 100644 --- a/xllm_ops/build_aclnn.sh +++ b/xllm_ops/build_aclnn.sh @@ -329,7 +329,7 @@ elif [[ "$SOC_VERSION" =~ ^ascend950 ]]; then # ### JD's in-house operators #### "beam_search_group" - # "x_attention" # A5 暂不参与编译, 集中解决 x_flash_attention_infer + "x_attention" "cache_unshared_kv" "causal_conv1d" "causal_conv1d_qkv" diff --git a/xllm_ops/x_attention/op_host/CMakeLists.txt b/xllm_ops/x_attention/op_host/CMakeLists.txt index 34ecd1f..b420975 100644 --- a/xllm_ops/x_attention/op_host/CMakeLists.txt +++ b/xllm_ops/x_attention/op_host/CMakeLists.txt @@ -15,15 +15,50 @@ if (BUILD_OPEN_PROJECT) ) endif() +# Dynamically set CATLASS_ARCH based on the SOC being built. +# NOTE: In the CMake scope SOC_VERSION may be empty; the reliable variable is +# ASCEND_COMPUTE_UNIT (see CMakeCache, e.g. "ascend950"). We accept both and also +# any *950 / *310p5 spelling. On A5 we inject -DCATLASS_ARCH=3510 so the arch35 +# (A5) branch is compiled in for both kernel and host; on A3 it is empty and the +# A3 (AtlasA2) branch is used. The host tiling object (ophost_xllm_tiling_obj) +# is built by the framework in another directory, so per-source compile defs from +# here cannot reach it; instead we generate xa_arch_config.h (included by the +# tiling sources) to bake the arch selection in at configure time. +string(TOLOWER "${SOC_VERSION}" _XA_SOC_LOWER) +string(TOLOWER "${ASCEND_COMPUTE_UNIT}" _XA_UNIT_LOWER) +if(_XA_SOC_LOWER MATCHES "ascend950" OR _XA_SOC_LOWER MATCHES "ascend310p5" + OR _XA_UNIT_LOWER MATCHES "ascend950" OR _XA_UNIT_LOWER MATCHES "ascend310p5") + set(CATLASS_ARCH_DEF "-DCATLASS_ARCH=3510") + set(_XA_IS_A5 TRUE) +else() + set(CATLASS_ARCH_DEF "") + set(_XA_IS_A5 FALSE) +endif() + add_ops_compile_options( OP_NAME XAttention OPTIONS --cce-auto-sync=on -Wno-deprecated-declarations -Werror + ${CATLASS_ARCH_DEF} -I${CANN_3RD_LIB_PATH}/catlass/include -I${CMAKE_CURRENT_LIST_DIR}/../../../ ) +# Generate per-build arch config header consumed by the host tiling sources +# (x_attention_tiling.h includes xa_arch_config.h). A5 builds get +# `#define CATLASS_ARCH 3510`; A3 builds get an empty header. +if(_XA_IS_A5) + set(XA_ARCH_CONFIG_BODY "#define CATLASS_ARCH 3510") +else() + set(XA_ARCH_CONFIG_BODY "") +endif() +configure_file( + ${CMAKE_CURRENT_SOURCE_DIR}/xa_arch_config.h.in + ${CMAKE_CURRENT_SOURCE_DIR}/xa_arch_config.h + @ONLY +) + if (NOT BUILD_OPS_RTY_KERNEL) add_modules_sources(OPTYPE x_attention ACLNNTYPE aclnn) endif() \ No newline at end of file diff --git a/xllm_ops/x_attention/op_host/x_attention_tiling.cpp b/xllm_ops/x_attention/op_host/x_attention_tiling.cpp index 572b948..7335ae8 100644 --- a/xllm_ops/x_attention/op_host/x_attention_tiling.cpp +++ b/xllm_ops/x_attention/op_host/x_attention_tiling.cpp @@ -41,13 +41,245 @@ constexpr int32_t NUM2 = 2; constexpr int32_t NUM3 = 3; constexpr int32_t NUM4 = 4; constexpr int32_t UNSHARED_Q_TILE = 128; +constexpr int32_t Q_S_BLOCK_TILE = 128; +constexpr int32_t BLOCK_SIZE = 128; +constexpr int32_t FLOAT_BLOCK_SIZE = 8; +constexpr int32_t SCALE_VALUE_ATTR_INDEX = 0; + +#if defined(CATLASS_ARCH) && (CATLASS_ARCH == 3510) +// ============================================================================ +// A5(Ascend950/DAV_3510) tiling implementation +// ============================================================================ +constexpr int32_t UNSHARED_KV_TILE = 128; +constexpr int32_t COMBINE_MAX_ROW_NUM_PER_LOOP = 64; + +class TilingXAttentionFunc { + public: + explicit TilingXAttentionFunc(gert::TilingContext* tiling_context) + : tiling_context_(tiling_context) {} + ge::graphStatus RunTiling(); + private: + uint64_t GetTilingKey() const; + private: + XAttentionTilingData tiling_data_; + gert::TilingContext* tiling_context_ = nullptr; + uint32_t sharedBlockDim = 0; + uint32_t unsharedBlockDim = 0; + uint32_t cubeCoreNum; + uint32_t vecCoreNum; + uint64_t ubSize; + ge::graphStatus FillBasicTilingData(); + void FillSharedSplitCoreTilingData(); + void FillUnsharedSplitCoreTilingData(); + void FillCombineScaleTilingData(); + void BalanceAicore(); + void SetWorkspaces(); +}; + + +void TilingXAttentionFunc::BalanceAicore() +{ + // Support dynamic calculation based on the amount of computation in the future. + sharedBlockDim = cubeCoreNum / 2; + unsharedBlockDim = cubeCoreNum - sharedBlockDim; + return; +} + +void TilingXAttentionFunc::FillUnsharedSplitCoreTilingData() +{ + auto kvBatchStride = tiling_data_.baseInfo.get_beamSize() * + tiling_data_.baseInfo.get_kvHeads() * + tiling_data_.baseInfo.get_maxDecodeStep() * + tiling_data_.baseInfo.get_headDim(); + tiling_data_.unsharedInfo.set_kvBatchStride(kvBatchStride); + + tiling_data_.unsharedInfo.set_usedCoreNum(unsharedBlockDim); + auto groupSize = tiling_data_.baseInfo.get_groupSize(); + uint32_t totalGroupCount = tiling_data_.baseInfo.get_beamSize() * tiling_data_.baseInfo.get_kvHeads(); + uint32_t maxGroupCountPerLoop = std::min(UNSHARED_Q_TILE / groupSize, UNSHARED_KV_TILE / tiling_data_.baseInfo.get_maxDecodeStep()); + while (maxGroupCountPerLoop > 1 && (totalGroupCount % maxGroupCountPerLoop != 0)) { + --maxGroupCountPerLoop; + } + tiling_data_.unsharedInfo.set_groupCountPerLoop(maxGroupCountPerLoop); + uint32_t perBatchTaskNum = totalGroupCount / maxGroupCountPerLoop; + uint32_t totalTaskNum = perBatchTaskNum * tiling_data_.baseInfo.get_batchSize(); + uint32_t perCoreTaskNum = (totalTaskNum + unsharedBlockDim - 1) / unsharedBlockDim; + tiling_data_.unsharedInfo.set_perBatchTaskNum(perBatchTaskNum); + tiling_data_.unsharedInfo.set_perCoreTaskNum(perCoreTaskNum); + tiling_data_.unsharedInfo.set_totalTaskNum(totalTaskNum); +} + +void TilingXAttentionFunc::SetWorkspaces() +{ + auto platform_info = + platform_ascendc::PlatformAscendC(tiling_context_->GetPlatformInfo()); + size_t systemWorkspaceSize = static_cast(platform_info.GetLibApiWorkSpaceSize()); + size_t userWorkspaceSize = 0; + + auto totalTokensQ = tiling_data_.baseInfo.get_totalTokensQ(); + auto qHeads = tiling_data_.baseInfo.get_qHeads(); + auto headDim = tiling_data_.baseInfo.get_headDim(); + uint64_t qOSize = totalTokensQ * qHeads * headDim * sizeof(float); + uint64_t sumMaxSize = totalTokensQ * qHeads * sizeof(float); + uint64_t sharedWorkspaceSize = qOSize + sumMaxSize * 2; + userWorkspaceSize = sharedWorkspaceSize * 2; + tiling_data_.set_qOSize(qOSize); + tiling_data_.set_sumMaxSize(sumMaxSize); + tiling_data_.set_sharedWorkspaceSize(sharedWorkspaceSize); + + size_t* workspace = tiling_context_->GetWorkspaceSizes(1); + workspace[0] = systemWorkspaceSize + userWorkspaceSize; +} + +void TilingXAttentionFunc::FillCombineScaleTilingData() +{ + auto totalTokensQ = tiling_data_.baseInfo.get_totalTokensQ(); + auto headDim = tiling_data_.baseInfo.get_headDim(); + auto qHeads = tiling_data_.baseInfo.get_qHeads(); + int32_t rowNum = totalTokensQ * qHeads; + int32_t bufferNum = 2; + // sharedMax、sharedSum、unsharedMax、unsharedSum + uint64_t maxReduceUbSize = COMBINE_MAX_ROW_NUM_PER_LOOP * sizeof(float) * (bufferNum * 4 + 3); + uint64_t remainUbSize = ubSize - maxReduceUbSize; + int32_t rowPerLoop = remainUbSize / (sizeof(float) * 3 * headDim * bufferNum); + if (rowPerLoop > COMBINE_MAX_ROW_NUM_PER_LOOP) { + rowPerLoop = COMBINE_MAX_ROW_NUM_PER_LOOP; + } + + int32_t totalTaskNum = (rowNum + rowPerLoop - 1) / rowPerLoop; + int32_t combineUsedCoreNum; + int32_t combineFormerCoreNum; + int32_t combineFormerTaskNum; + int32_t combineTailTaskNum; + + if (totalTaskNum <= vecCoreNum) { + combineUsedCoreNum = totalTaskNum; + combineFormerCoreNum = 0; + combineFormerTaskNum = 1; + combineTailTaskNum = 1; + } else { + combineUsedCoreNum = vecCoreNum; + int32_t taskNumPerCore = totalTaskNum / combineUsedCoreNum; + int32_t taskNumTailPerCore = totalTaskNum % combineUsedCoreNum; + combineFormerCoreNum = taskNumTailPerCore; + combineFormerTaskNum = taskNumPerCore + 1; + combineTailTaskNum = taskNumPerCore; + } + + tiling_data_.combineInfo.set_rowPerLoop(rowPerLoop); + tiling_data_.combineInfo.set_rowNum(rowNum); + tiling_data_.combineInfo.set_totalTaskNum(totalTaskNum); + tiling_data_.combineInfo.set_formerCoreNum(combineFormerCoreNum); + tiling_data_.combineInfo.set_formerTaskNum(combineFormerTaskNum); + tiling_data_.combineInfo.set_tailTaskNum(combineTailTaskNum); + tiling_data_.combineInfo.set_usedCoreNum(combineUsedCoreNum); +} + +void TilingXAttentionFunc::FillSharedSplitCoreTilingData() +{ + tiling_data_.sharedInfo.set_usedCoreNum(sharedBlockDim); + auto beamSize = tiling_data_.baseInfo.get_beamSize(); + auto qHeads = tiling_data_.baseInfo.get_qHeads(); + auto batchSize = tiling_data_.baseInfo.get_batchSize(); + int32_t perBatchHeadTaskNum = (beamSize + Q_S_BLOCK_TILE - 1) / Q_S_BLOCK_TILE; + int32_t totalTaskNum = perBatchHeadTaskNum * qHeads * batchSize; + int32_t perCoreTaskNum = (totalTaskNum + sharedBlockDim - 1) / sharedBlockDim; + tiling_data_.sharedInfo.set_totalTaskNum(totalTaskNum); + tiling_data_.sharedInfo.set_perBatchHeadTaskNum(perBatchHeadTaskNum); + tiling_data_.sharedInfo.set_perCoreTaskNum(perCoreTaskNum); +} + +ge::graphStatus TilingXAttentionFunc::FillBasicTilingData() +{ + auto sharedBlockTableShapePtr = tiling_context_->GetOptionalInputShape(InputIndex::SHARED_BLOCK_TABLE); + auto unsharedBlockTableShapePtr = tiling_context_->GetOptionalInputShape(InputIndex::UNSHARED_BLOCK_TABLE); + bool isSharedPaged = (sharedBlockTableShapePtr != nullptr); + bool isUnsharedPaged = (unsharedBlockTableShapePtr != nullptr); + + if (!(!isSharedPaged && isUnsharedPaged)) { + OP_LOGE(tiling_context_->GetNodeName(), "xAttention only support unshared_paged and not shared_paged on Ascend950."); + return ge::GRAPH_FAILED; + } + + auto queryShape = tiling_context_->GetInputShape(QUERY)->GetStorageShape(); + auto sharedKeyBlockShape = tiling_context_->GetInputShape(SHARED_KEY_BLOCK)->GetStorageShape(); + auto unsharedKeyBlockShape = tiling_context_->GetInputShape(UNSHARED_KEY_BLOCK)->GetStorageShape(); + auto unsharedBlockTableShape = tiling_context_->GetOptionalInputShape(UNSHARED_BLOCK_TABLE)->GetStorageShape(); + auto sharedKvLenShape = tiling_context_->GetInputShape(SHARED_KV_LENS)->GetStorageShape(); + + + int32_t totalTokensQ = queryShape.GetDim(0); + int32_t sharedKvTokens = sharedKeyBlockShape.GetDim(0); + int32_t qHeads = queryShape.GetDim(1); + int32_t kvHeads = sharedKeyBlockShape.GetDim(1); + int32_t headDim = queryShape.GetDim(2); + int32_t beamSize = unsharedKeyBlockShape.GetDim(1); + int32_t maxDecodeStep = unsharedKeyBlockShape.GetDim(3); + int32_t batchSize = sharedKvLenShape.GetDim(0); + int32_t groupSize = qHeads / kvHeads; + + float scaleValue = static_cast(1.0 / std::sqrt(1.0 * headDim)); + auto attrs = tiling_context_->GetAttrs(); + if (attrs != nullptr) { + const auto* attr_scale_value = attrs->GetAttrPointer(SCALE_VALUE_ATTR_INDEX); + if (attr_scale_value != nullptr && *attr_scale_value > 0.0f) { + scaleValue = *attr_scale_value; + } + } + + tiling_data_.baseInfo.set_batchSize(batchSize); + tiling_data_.baseInfo.set_beamSize(beamSize); + tiling_data_.baseInfo.set_qHeads(qHeads); + tiling_data_.baseInfo.set_kvHeads(kvHeads); + tiling_data_.baseInfo.set_groupSize(groupSize); + tiling_data_.baseInfo.set_headDim(headDim); + tiling_data_.baseInfo.set_scaleValue(scaleValue); + tiling_data_.baseInfo.set_totalTokensQ(totalTokensQ); + tiling_data_.baseInfo.set_sharedKvTokens(sharedKvTokens); + tiling_data_.baseInfo.set_maxDecodeStep(maxDecodeStep); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingXAttentionFunc::RunTiling() +{ + // Get platform hardware information + auto platformInfo = + platform_ascendc::PlatformAscendC(tiling_context_->GetPlatformInfo()); + cubeCoreNum = platformInfo.GetCoreNumAic(); + vecCoreNum = platformInfo.GetCoreNumAiv(); + platformInfo.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); + + BalanceAicore(); + auto ret = FillBasicTilingData(); + if (ret != ge::GRAPH_SUCCESS) { + OP_LOGE(tiling_context_->GetNodeName(), "fill basic tiling failed."); + return ge::GRAPH_FAILED; + } + + FillSharedSplitCoreTilingData(); + FillUnsharedSplitCoreTilingData(); + FillCombineScaleTilingData(); + SetWorkspaces(); + + // Save tilingData + tiling_data_.SaveToBuffer(tiling_context_->GetRawTilingData()->GetData(), + tiling_context_->GetRawTilingData()->GetCapacity()); + tiling_context_->GetRawTilingData()->SetDataSize(tiling_data_.GetDataSize()); + tiling_context_->SetBlockDim(cubeCoreNum); + + // tiling_context_->SetTilingKey(GetTilingKey()); + + return ge::GRAPH_SUCCESS; +} + +#else +// ============================================================================ +// A3(AtlasA2/A3) tiling implementation +// ============================================================================ constexpr int32_t UNSHARED_KV_TILE = 256; -constexpr uint32_t Q_S_BLOCK_TILE = 128; -constexpr uint32_t BLOCK_SIZE = 128; +constexpr uint32_t Q_S_BLOCK_TILE_A3 = 128; constexpr int32_t WORKSPACE_BLOCK_SIZE_DB = 128 * 128 * 4; // row * col * blockStackNum constexpr int32_t UNSHARED_WORKSPACE_BLOCK_SIZE_DB = 128 * 256; // unshared no pinpong -constexpr int32_t FLOAT_BLOCK_SIZE = 8; -constexpr int32_t SCALE_VALUE_ATTR_INDEX = 0; class TilingXAttentionFunc { public: @@ -75,7 +307,7 @@ class TilingXAttentionFunc { void BalanceAicore(); void SetWorkspaces(); uint32_t GetQNBlockTile(int64_t qSeqlen, uint32_t groupSize); - + }; @@ -88,7 +320,7 @@ ge::graphStatus TilingXAttentionFunc::FillBasicTilingData4NewKind() // unshared_blk_tb [bs, request_idx] auto unsharedKeyBlockShape = tiling_context_->GetInputShape(UNSHARED_KEY_BLOCK)->GetStorageShape(); auto unsharedBlockTableShape = tiling_context_->GetOptionalInputShape(UNSHARED_BLOCK_TABLE)->GetStorageShape(); - + int32_t numTokens = queryShape.GetDim(0); int32_t qHeadNum = queryShape.GetDim(1); int32_t embeddingSize = queryShape.GetDim(2); @@ -208,14 +440,14 @@ void TilingXAttentionFunc::SetWorkspaces() platform_ascendc::PlatformAscendC(tiling_context_->GetPlatformInfo()); size_t systemWorkspaceSize = static_cast(platform_info.GetLibApiWorkSpaceSize()); size_t userWorkspaceSize = 0; - - uint64_t qoSize = tiling_data_.get_numTokens() + + uint64_t qoSize = tiling_data_.get_numTokens() * tiling_data_.get_numHeads() * tiling_data_.get_embeddingSize() * sizeof(int16_t); // Attention occupied space // TODO: Only apply for one temporary space, affecting preload function, long sequence scenario needs extra processing - uint64_t mm1OutSize = (sharedBlockDim * WORKSPACE_BLOCK_SIZE_DB + + uint64_t mm1OutSize = (sharedBlockDim * WORKSPACE_BLOCK_SIZE_DB + unsharedBlockDim * UNSHARED_WORKSPACE_BLOCK_SIZE_DB) * NUM3 * sizeof(float);; uint64_t smOnlineOutSize = (sharedBlockDim * WORKSPACE_BLOCK_SIZE_DB + unsharedBlockDim * UNSHARED_WORKSPACE_BLOCK_SIZE_DB) * NUM3 * sizeof(int16_t); @@ -242,11 +474,11 @@ void TilingXAttentionFunc::SetWorkspaces() void TilingXAttentionFunc::FillCombineScaleTilingData() { - uint32_t rowNum = tiling_data_.get_batch() * - tiling_data_.get_beamSize() * + uint32_t rowNum = tiling_data_.get_batch() * + tiling_data_.get_beamSize() * tiling_data_.get_numHeads(); uint32_t columnSize = tiling_data_.get_embeddingSize(); - + uint32_t rowNumPerCore = rowNum / cubeCoreNum; // number of rows per core uint32_t rowNumTailPerCore = rowNum % cubeCoreNum; // remaining rows, need to be allocated to the first few cores tiling_data_.set_combineFormerCoreNum(rowNumTailPerCore); @@ -263,7 +495,7 @@ void TilingXAttentionFunc::FillSharedSplitCoreTilingData() uint32_t curQNBlockTile = GetQNBlockTile(qSeqlen, groupSize); uint32_t qNBlockNumPerGroup = (groupSize + curQNBlockTile - 1) / curQNBlockTile; uint32_t curQNBlockNum = qNBlockNumPerGroup * tiling_data_.get_kvHeads(); - uint32_t curQSBlockTile = Q_S_BLOCK_TILE; + uint32_t curQSBlockTile = Q_S_BLOCK_TILE_A3; uint32_t curQSBlockNum = (qSeqlen + curQSBlockTile - 1) / curQSBlockTile; uint32_t curTaskNum = curQNBlockNum * curQSBlockNum; uint32_t firstSharedBatchTaskNum = curTaskNum; @@ -280,7 +512,7 @@ ge::graphStatus TilingXAttentionFunc::FillBasicTilingData() auto sharedKeyBlockShape = tiling_context_->GetInputShape(SHARED_KEY_BLOCK)->GetStorageShape(); auto unsharedKeyBlockShape = tiling_context_->GetInputShape(UNSHARED_KEY_BLOCK)->GetStorageShape(); auto sharedBlockTableShape = tiling_context_->GetOptionalInputShape(SHARED_BLOCK_TABLE)->GetStorageShape(); - + int32_t numTokens = queryShape.GetDim(0); int32_t qHeadNum = queryShape.GetDim(1); int32_t embeddingSize = queryShape.GetDim(2); @@ -355,6 +587,7 @@ ge::graphStatus TilingXAttentionFunc::RunTiling() return ge::GRAPH_SUCCESS; } +#endif static ge::graphStatus TilingFunc(gert::TilingContext* context) diff --git a/xllm_ops/x_attention/op_host/x_attention_tiling.h b/xllm_ops/x_attention/op_host/x_attention_tiling.h index 4cb92c3..cb10275 100644 --- a/xllm_ops/x_attention/op_host/x_attention_tiling.h +++ b/xllm_ops/x_attention/op_host/x_attention_tiling.h @@ -13,10 +13,75 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +// Per-build arch selection (generated by CMake). On A5 this defines +// CATLASS_ARCH=3510 so the arch35 (A5) host tiling struct/branch is compiled in; +// on A3 it is empty. Included here so EVERY translation unit that pulls in this +// header (tiling.cpp, proto.cpp, ...) sees a consistent XAttentionTilingData +// layout, avoiding ODR violations. +#include "xa_arch_config.h" #include "register/tilingdata_base.h" #include "tiling/tiling_api.h" namespace optiling { + +#if defined(CATLASS_ARCH) && (CATLASS_ARCH == 3510) +// ===== A5(Ascend950/DAV_3510): nested tiling data ===== +BEGIN_TILING_DATA_DEF(XABaseInfoTilingData) + TILING_DATA_FIELD_DEF(int32_t, batchSize); + TILING_DATA_FIELD_DEF(int32_t, beamSize); + TILING_DATA_FIELD_DEF(int32_t, qHeads); + TILING_DATA_FIELD_DEF(int32_t, kvHeads); + TILING_DATA_FIELD_DEF(int32_t, groupSize); + TILING_DATA_FIELD_DEF(int32_t, headDim); + TILING_DATA_FIELD_DEF(float, scaleValue); + TILING_DATA_FIELD_DEF(int32_t, totalTokensQ); + TILING_DATA_FIELD_DEF(int32_t, sharedKvTokens); + TILING_DATA_FIELD_DEF(int32_t, maxDecodeStep); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(XABaseInfoTilingDataOp, XABaseInfoTilingData) + +BEGIN_TILING_DATA_DEF(XASharedTilingData) + TILING_DATA_FIELD_DEF(int32_t, totalTaskNum); + TILING_DATA_FIELD_DEF(int32_t, perBatchHeadTaskNum); + TILING_DATA_FIELD_DEF(int32_t, perCoreTaskNum); + TILING_DATA_FIELD_DEF(int32_t, usedCoreNum); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(XASharedTilingDataOp, XASharedTilingData) + +BEGIN_TILING_DATA_DEF(XAUnsharedTilingData) + TILING_DATA_FIELD_DEF(int32_t, kvBatchStride); + TILING_DATA_FIELD_DEF(int32_t, groupCountPerLoop); + TILING_DATA_FIELD_DEF(int32_t, perBatchTaskNum); + TILING_DATA_FIELD_DEF(int32_t, perCoreTaskNum); + TILING_DATA_FIELD_DEF(int32_t, totalTaskNum); + TILING_DATA_FIELD_DEF(int32_t, usedCoreNum); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(XAUnsharedTilingDataOp, XAUnsharedTilingData) + +BEGIN_TILING_DATA_DEF(XACombineTilingData) + TILING_DATA_FIELD_DEF(int32_t, rowPerLoop); + TILING_DATA_FIELD_DEF(int32_t, rowNum); + TILING_DATA_FIELD_DEF(int32_t, totalTaskNum); + TILING_DATA_FIELD_DEF(int32_t, formerCoreNum); + TILING_DATA_FIELD_DEF(int32_t, formerTaskNum); + TILING_DATA_FIELD_DEF(int32_t, tailTaskNum); + TILING_DATA_FIELD_DEF(int32_t, usedCoreNum); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(XACombineTilingDataOp, XACombineTilingData) + +BEGIN_TILING_DATA_DEF(XAttentionTilingData) + TILING_DATA_FIELD_DEF(uint64_t, qOSize); + TILING_DATA_FIELD_DEF(uint64_t, sumMaxSize); + TILING_DATA_FIELD_DEF(uint64_t, sharedWorkspaceSize); + TILING_DATA_FIELD_DEF_STRUCT(XABaseInfoTilingData, baseInfo); + TILING_DATA_FIELD_DEF_STRUCT(XASharedTilingData, sharedInfo); + TILING_DATA_FIELD_DEF_STRUCT(XAUnsharedTilingData, unsharedInfo); + TILING_DATA_FIELD_DEF_STRUCT(XACombineTilingData, combineInfo); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(XAttention, XAttentionTilingData) +#else +// ===== A3(AtlasA2/A3): flat tiling data ===== BEGIN_TILING_DATA_DEF(XAttentionTilingData) TILING_DATA_FIELD_DEF(uint32_t, numHeads); TILING_DATA_FIELD_DEF(uint32_t, kvHeads); @@ -53,4 +118,6 @@ BEGIN_TILING_DATA_DEF(XAttentionTilingData) END_TILING_DATA_DEF; REGISTER_TILING_DATA_CLASS(XAttention, XAttentionTilingData) +#endif + } diff --git a/xllm_ops/x_attention/op_host/xa_arch_config.h b/xllm_ops/x_attention/op_host/xa_arch_config.h new file mode 100644 index 0000000..f2e6b7c --- /dev/null +++ b/xllm_ops/x_attention/op_host/xa_arch_config.h @@ -0,0 +1,21 @@ +// ----------------------------------------------------------------------------------------------------------- +// Copyright (c) 2025 Huawei Technologies Co., Ltd. +// This file is generated by CMake (configure_file) from xa_arch_config.h.in. +// DO NOT EDIT the generated header directly. +// +// Purpose: The host tiling.cpp is compiled into a framework-aggregated OBJECT +// library (ophost_xllm_tiling_obj) defined in another directory, so per-source +// COMPILE_DEFINITIONS / add_compile_definitions from this op subdir cannot reach +// it. Instead we bake the A5 (arch35) selection into a generated header that the +// tiling sources include. The value is decided at CMake configure time based on +// ASCEND_COMPUTE_UNIT / SOC_VERSION, so A3 and A5 builds get different content. +// ----------------------------------------------------------------------------------------------------------- +#ifndef XA_ARCH_CONFIG_H +#define XA_ARCH_CONFIG_H + +// #define CATLASS_ARCH 3510 is replaced by CMake: +// - on A5 (ascend950 / ascend310p5): "#define CATLASS_ARCH 3510" +// - otherwise : (empty) +#define CATLASS_ARCH 3510 + +#endif // XA_ARCH_CONFIG_H diff --git a/xllm_ops/x_attention/op_host/xa_arch_config.h.in b/xllm_ops/x_attention/op_host/xa_arch_config.h.in new file mode 100644 index 0000000..51a72bb --- /dev/null +++ b/xllm_ops/x_attention/op_host/xa_arch_config.h.in @@ -0,0 +1,21 @@ +// ----------------------------------------------------------------------------------------------------------- +// Copyright (c) 2025 Huawei Technologies Co., Ltd. +// This file is generated by CMake (configure_file) from xa_arch_config.h.in. +// DO NOT EDIT the generated header directly. +// +// Purpose: The host tiling.cpp is compiled into a framework-aggregated OBJECT +// library (ophost_xllm_tiling_obj) defined in another directory, so per-source +// COMPILE_DEFINITIONS / add_compile_definitions from this op subdir cannot reach +// it. Instead we bake the A5 (arch35) selection into a generated header that the +// tiling sources include. The value is decided at CMake configure time based on +// ASCEND_COMPUTE_UNIT / SOC_VERSION, so A3 and A5 builds get different content. +// ----------------------------------------------------------------------------------------------------------- +#ifndef XA_ARCH_CONFIG_H +#define XA_ARCH_CONFIG_H + +// @XA_ARCH_CONFIG_BODY@ is replaced by CMake: +// - on A5 (ascend950 / ascend310p5): "#define CATLASS_ARCH 3510" +// - otherwise : (empty) +@XA_ARCH_CONFIG_BODY@ + +#endif // XA_ARCH_CONFIG_H diff --git a/xllm_ops/x_attention/op_kernel/arch35/combine_kernel.h b/xllm_ops/x_attention/op_kernel/arch35/combine_kernel.h new file mode 100644 index 0000000..bfc36ac --- /dev/null +++ b/xllm_ops/x_attention/op_kernel/arch35/combine_kernel.h @@ -0,0 +1,150 @@ +/* Copyright 2026 The xLLM Authors. 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 + + https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE + +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 "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 "catlass/epilogue/dispatch_policy.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" +#include "x_attention_common.h" +#include "kernel_operator.h" + +template +class CombineScaleKernel { +public: + using ArchTag = typename EpilogueCombineScale::ArchTag; + using ElementOutput = typename EpilogueCombineScale::ElementOutput; + using ElementInput = typename EpilogueCombineScale::ElementInput; + + CATLASS_DEVICE + CombineScaleKernel(XAttentionTilingData* tilingDataPtr): faTilingData(tilingDataPtr) {} + + template + CATLASS_DEVICE void operator()(XAttnKernelCommonParams const ¶ms); + + template <> + CATLASS_DEVICE void operator()(XAttnKernelCommonParams const ¶ms) { + return; + } + + template <> + CATLASS_DEVICE void operator()(XAttnKernelCommonParams const ¶ms) { + uint32_t ubBufAddrStart = 0; + uint32_t rowNumPerLoop = faTilingData->combineInfo.rowPerLoop; + uint32_t headDim = faTilingData->baseInfo.headDim; + EpilogueCombineScale epilogueCombineScale(resource, ubBufAddrStart, rowNumPerLoop, headDim); + + int32_t usedCoreNum = faTilingData->combineInfo.usedCoreNum; + int32_t coreIdx = AscendC::GetBlockIdx(); + + if (coreIdx >= usedCoreNum) { + return; + } + + int32_t formerCoreNum = faTilingData->combineInfo.formerCoreNum; + int32_t tailCoreNum = usedCoreNum - formerCoreNum; + int32_t totalTaskNum = faTilingData->combineInfo.totalTaskNum; + int32_t formerTaskNum = faTilingData->combineInfo.formerTaskNum; + int32_t tailTaskNum = faTilingData->combineInfo.tailTaskNum; + int32_t rowNum = faTilingData->combineInfo.rowNum; + int32_t coreTaskNum; + int32_t mainTaskRowNum; + int32_t tailTaskRowNum; + int32_t attnOffsetPerCore; + int32_t gmglOffsetPerCore; + + // AscendC::printf("rowNum %d rowNumPerLoop %d formerCoreNum %d formerTaskNum %d tailTaskNum %d usedCoreNum %d tailCoreNum %d\n", rowNum, + // rowNumPerLoop, formerCoreNum, formerTaskNum, tailTaskNum, usedCoreNum, tailCoreNum); + + if (coreIdx < formerCoreNum) { + coreTaskNum = formerTaskNum; + mainTaskRowNum = rowNumPerLoop; + tailTaskRowNum = rowNumPerLoop; + gmglOffsetPerCore = coreIdx * rowNumPerLoop * coreTaskNum; + attnOffsetPerCore = gmglOffsetPerCore * headDim; + } else { + coreTaskNum = tailTaskNum; + mainTaskRowNum = rowNumPerLoop; + tailTaskRowNum = rowNum - formerCoreNum * formerTaskNum * rowNumPerLoop - (tailCoreNum - 1) * mainTaskRowNum; + gmglOffsetPerCore = (formerCoreNum * formerTaskNum + (coreIdx - formerCoreNum) * tailTaskNum) * rowNumPerLoop; + attnOffsetPerCore = gmglOffsetPerCore * headDim; + } + + AscendC::GlobalTensor gSharedGm; + gSharedGm.SetGlobalBuffer((__gm__ ElementInput *)params.sharedMax + gmglOffsetPerCore); + AscendC::GlobalTensor gSharedGl; + gSharedGl.SetGlobalBuffer((__gm__ ElementInput *)params.sharedSum + gmglOffsetPerCore); + AscendC::GlobalTensor gUnsharedGm; + gUnsharedGm.SetGlobalBuffer((__gm__ ElementInput *)params.unsharedMax + gmglOffsetPerCore); + AscendC::GlobalTensor gUnsharedGl; + gUnsharedGl.SetGlobalBuffer((__gm__ ElementInput *)params.unsharedSum + gmglOffsetPerCore); + AscendC::GlobalTensor gSharedOut; + gSharedOut.SetGlobalBuffer((__gm__ ElementInput *)params.sharedO + attnOffsetPerCore); + AscendC::GlobalTensor gUnsharedOut; + gUnsharedOut.SetGlobalBuffer((__gm__ ElementInput *)params.unsharedO + attnOffsetPerCore); + AscendC::GlobalTensor gFinalOut; + gFinalOut.SetGlobalBuffer((__gm__ ElementOutput *)params.o + attnOffsetPerCore); + + // if (coreIdx == 0) { + // for (int i = 68; i < 69; i++) { + // AscendC::printf("token %d sharedOut\n", i); + // AscendC::DumpTensor(gSharedOut[i * headDim], 1, 8); + // AscendC::printf("token %d sharedMax %f\n", i, gSharedGm.GetValue(i)); + // AscendC::printf("token %d sharedSum %f\n", i, gSharedGl.GetValue(i)); + // AscendC::printf("token %d unsharedOut\n", i); + // AscendC::DumpTensor(gUnsharedOut[i * headDim], 3, 8); + // AscendC::printf("token %d unsharedMax %f\n", i, gUnsharedGm.GetValue(i)); + // AscendC::printf("token %d unsharedSum %f\n", i, gUnsharedGl.GetValue(i)); + // } + // } + + + int8_t taskId = 0; + for (int i = 0; i < coreTaskNum; i++) { + // int32_t realRowNum = (i == coreTaskNum - 1) ? tailTaskRowNum : mainTaskRowNum; + int64_t gmglTaskOffset = i * rowNumPerLoop; + int64_t globalRowStart = gmglOffsetPerCore + gmglTaskOffset; + int32_t remainingRows = rowNum - globalRowStart; + int32_t realRowNum = + remainingRows < static_cast(rowNumPerLoop) + ? remainingRows + : static_cast(rowNumPerLoop); + + if (realRowNum <= 0) { + break; + } + int64_t attnTaskOffset = gmglTaskOffset * headDim; + epilogueCombineScale( + gSharedGm[gmglTaskOffset], + gUnsharedGm[gmglTaskOffset], + gSharedGl[gmglTaskOffset], + gUnsharedGl[gmglTaskOffset], + gSharedOut[attnTaskOffset], + gUnsharedOut[attnTaskOffset], + gFinalOut[attnTaskOffset], + realRowNum, + taskId + ); + } + } +private: + Arch::Resource resource; + XAttentionTilingData* faTilingData; +}; diff --git a/xllm_ops/x_attention/op_kernel/arch35/shared_infer_catlass_kernel.h b/xllm_ops/x_attention/op_kernel/arch35/shared_infer_catlass_kernel.h new file mode 100644 index 0000000..9788dd6 --- /dev/null +++ b/xllm_ops/x_attention/op_kernel/arch35/shared_infer_catlass_kernel.h @@ -0,0 +1,444 @@ + +#ifndef X_ATTN_SHARED_FA_INFER_CATLASS_KERNEL_H +#define X_ATTN_SHARED_FA_INFER_CATLASS_KERNEL_H + +#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 "catlass/epilogue/dispatch_policy.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" +#include "x_attention_common.h" +#include "kernel_operator.h" + +using namespace Catlass; + +template < + class BlockMmadQK, + class BlockMmadPV, + class EpilogueFASoftmax, + class EpilogueFARescale, + typename KVLEN_T> +class SharedFaInferKernel { + public: + using ArchTag = typename BlockMmadQK::ArchTag; + using L1TileShape = typename BlockMmadQK::L1TileShape; + using ElementQ = typename BlockMmadQK::ElementA; + using LayoutTagQ = typename BlockMmadQK::LayoutTagA; + using ElementK = typename BlockMmadQK::ElementB; + using LayoutTagK = typename BlockMmadQK::LayoutTagB; + using ElementS = typename BlockMmadQK::ElementC; + using LayoutTagS = typename BlockMmadQK::LayoutTagC; + + using ElementP = typename BlockMmadPV::ElementA; + using LayoutTagP = typename BlockMmadPV::LayoutTagA; + using LayoutTagPL1 = typename BlockMmadPV::TileCopy::LayoutTagL1A; + using ElementV = typename BlockMmadPV::ElementB; + using LayoutTagV = typename BlockMmadPV::LayoutTagB; + using ElementOTmp = typename BlockMmadPV::ElementC; + using LayoutTagOTmp = typename BlockMmadPV::LayoutTagC; + + static constexpr uint32_t qSeqlenTemplateType = tla::get<0>(L1TileShape{}); + static constexpr uint32_t kvSeqlenTemplateType = tla::get<1>(L1TileShape{}); + static constexpr uint32_t embedTemplateType = tla::get<2>(L1TileShape{}); + static constexpr uint32_t halfQSeqlenTemplateType = qSeqlenTemplateType / CV_RATIO; + + CATLASS_DEVICE + SharedFaInferKernel(XAttentionTilingData *tilingData) { + this->tilingData = tilingData; + } + + template + CATLASS_DEVICE void operator()(XAttnKernelCommonParams const ¶ms); + + CATLASS_DEVICE void Init(XAttnKernelCommonParams const ¶ms) { + auto qkSize = halfQSeqlenTemplateType * kvSeqlenTemplateType * sizeof(ElementS); + auto pvSize = halfQSeqlenTemplateType * embedTemplateType * sizeof(ElementOTmp); + + for (int i = 0; i < 2; i++) { + qkTensorList[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += qkSize; + pvTensorList[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += pvSize; + } + + auto reduceSize = halfQSeqlenTemplateType * sizeof(ElementS); + + if ASCEND_IS_AIV { + for (int i = 0; i < 3; i++) { + expSumUb[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += reduceSize; + expMaxUb[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += reduceSize; + maxUb[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += reduceSize; + } + } + + auto pL1Size = qSeqlenTemplateType * kvSeqlenTemplateType * sizeof(ElementP); + for (int i = 0; i < 3; i++) { + pL1TensorList[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); + l1BufAddrStart += pL1Size; + } + + sharedKvLensGm.SetGlobalBuffer((__gm__ KVLEN_T *)params.sharedKvLens); + + batchSize = tilingData->baseInfo.batchSize; + beamSize = tilingData->baseInfo.beamSize; + qHeads = tilingData->baseInfo.qHeads; + kvHeads = tilingData->baseInfo.kvHeads; + groupSize = tilingData->baseInfo.groupSize; + headDim = tilingData->baseInfo.headDim; + scaleValue = tilingData->baseInfo.scaleValue; + totalTokensQ = tilingData->baseInfo.totalTokensQ; + sharedKvTokens = tilingData->baseInfo.sharedKvTokens; + + coreNum = tilingData->sharedInfo.usedCoreNum; + coreIdx = AscendC::GetBlockIdx(); + + if ASCEND_IS_AIV { + coreIdx = coreIdx / CV_RATIO; + subVecIdx = AscendC::GetSubBlockIdx(); + } + + strideQO = qHeads * headDim; + strideKV = kvHeads * headDim; + } + + CATLASS_DEVICE void operator()(XAttnKernelCommonParams const ¶ms) { + uint32_t taskIdL0A = 0; + uint32_t taskIdL0B = 0; + + Init(params); + SetFlag(); + + AscendC::GlobalTensor gQ; + gQ.SetGlobalBuffer((__gm__ ElementQ *)params.q); + auto layoutQ = tla::MakeLayout(totalTokensQ, qHeads * headDim); + auto tensorQ = tla::MakeTensor(gQ, layoutQ, Arch::PositionGM{}); + AscendC::GlobalTensor gK; + gK.SetGlobalBuffer((__gm__ ElementK *)params.sharedK); + auto layoutK = tla::MakeLayout(kvHeads * headDim, sharedKvTokens); + auto tensorK = tla::MakeTensor(gK, layoutK, Arch::PositionGM{}); + AscendC::GlobalTensor gV; + gV.SetGlobalBuffer((__gm__ ElementV *)params.sharedV); + auto layoutV = tla::MakeLayout(sharedKvTokens, kvHeads * headDim); + auto tensorV = tla::MakeTensor(gV, layoutV, Arch::PositionGM{}); + + AscendC::GlobalTensor gSharedO; + gSharedO.SetGlobalBuffer((__gm__ ElementOTmp *)params.sharedO); + auto layoutO = tla::MakeLayout(totalTokensQ, qHeads * headDim); + auto tensorO = tla::MakeTensor(gSharedO, layoutO, Arch::PositionGM{}); + AscendC::GlobalTensor sharedMaxGm; + sharedMaxGm.SetGlobalBuffer((__gm__ ElementOTmp *)params.sharedMax); + AscendC::GlobalTensor sharedSumGm; + sharedSumGm.SetGlobalBuffer((__gm__ ElementOTmp *)params.sharedSum); + + BlockMmadQK blockMmadQK(resource, l1BufAddrStart, l0CBufAddrStart); + BlockMmadPV blockMmadPV(resource, l1BufAddrStart, l0CBufAddrStart); + EpilogueFASoftmax epilogueSoftmax(resource, ubBufAddrStart, scaleValue, qHeads); + EpilogueFARescale epilogueRescale(resource, ubBufAddrStart); + + int32_t taskId = 0; + int32_t perCoreTaskNum = tilingData->sharedInfo.perCoreTaskNum; + int32_t totalTaskNum = tilingData->sharedInfo.totalTaskNum; + + int32_t taskStartId = coreIdx * perCoreTaskNum; + if (taskStartId >= totalTaskNum) { + WaitFlag(); + return; + } + int32_t coreTaskNum = perCoreTaskNum; + int32_t taskEndId = taskStartId + coreTaskNum; + + if (taskEndId > totalTaskNum) { + taskEndId = totalTaskNum; + coreTaskNum = taskEndId - taskStartId; + } + + // AscendC::printf("coreNum %d coreIdx %d coreTaskNum %d taskStartId %d taskEndId %d\n", + // coreNum, coreIdx, coreTaskNum, taskStartId, taskEndId); + + for (int32_t qTaskId = taskStartId; qTaskId < taskEndId + 3; qTaskId++) + { + bool notLastThreeLoop = qTaskId < taskEndId; + bool notLastTwoLoop = qTaskId < taskEndId + 1; + bool notLast = qTaskId < taskEndId + 2; + + int32_t kvLen = 0; + int32_t kvBlockNum = 1; + + if (notLastThreeLoop) { + SharedInfer::TaskArgs taskArgs; + GetQTaskInfo(taskArgs, qTaskId); + taskArgList[taskId % 4] = taskArgs; + kvLen = taskArgs.actualKvLen; + kvBlockNum = (kvLen + kvSeqlenTemplateType - 1) / kvSeqlenTemplateType; + } + + for (int32_t kvBlockId = 0; kvBlockId < kvBlockNum; kvBlockId++) { + if (notLastThreeLoop) { + auto nowTaskId = taskId % 4; + SharedInfer::TaskArgs &taskArgsNow = taskArgList[nowTaskId]; + GetKvTaskInfo(taskArgsNow, kvBlockId, kvBlockNum, taskId); + + if ASCEND_IS_AIC { + auto actualShape = tla::MakeShape(taskArgsNow.blockQLen, taskArgsNow.blockKvLen, headDim); + auto tensorQTile = GetTile( + tensorQ, + tla::MakeCoord(taskArgsNow.qCoord, taskArgsNow.qNCoord), + tla::MakeShape(taskArgsNow.blockQLen, headDim) + ); + + auto tensorKTile = GetTile( + tensorK, + tla::MakeCoord(taskArgsNow.kvNCoord, taskArgsNow.kvCoord), + tla::MakeShape(headDim, taskArgsNow.blockKvLen) + ); + + auto layoutQKRes = tla::MakeLayout(taskArgsNow.blockQLen, kvSeqlenTemplateType); + auto tensorQKRes = tla::MakeTensor(qkTensorList[taskArgsNow.taskIdMod2], layoutQKRes, Arch::PositionUB{}); + blockMmadQK( + tensorQTile, tensorKTile, tensorQKRes, actualShape, + SharedInfer::QK_UB_RELEASE_FLAG[taskArgsNow.taskIdMod2], + taskArgsNow.isFirstKv, taskArgsNow.isLastKv, + taskIdL0A, taskIdL0B + ); + + AscendC::CrossCoreSetFlag(SharedInfer::SYNC_QK_READY_FLAG[taskArgsNow.taskIdMod2]); + AscendC::CrossCoreSetFlag(16 + SharedInfer::SYNC_QK_READY_FLAG[taskArgsNow.taskIdMod2]); + } + } + + if (taskId > 0 && notLastTwoLoop) { + if ASCEND_IS_AIV { + auto &taskArgsPre = taskArgList[(taskId - 1) % 4]; + auto qkResLayout = tla::MakeLayout(taskArgsPre.halfBlockQLen, taskArgsPre.blockKvLen); + auto qkResTensor = tla::MakeTensor(qkTensorList[taskArgsPre.taskIdMod2], qkResLayout, Arch::PositionUB{}); + auto pL1OutLayout = tla::MakeLayout(qSeqlenTemplateType, kvSeqlenTemplateType); + auto pL1OutTensor = tla::MakeTensor(pL1TensorList[taskArgsPre.taskIdMod3], pL1OutLayout, Arch::PositionL1{}); + auto pL1OutTile = GetTile( + pL1OutTensor, + tla::MakeCoord(taskArgsPre.halfBlockQOffset, 0), + tla::MakeShape(taskArgsPre.halfBlockQLen, kvSeqlenTemplateType) + ); + auto sharedMaxTile = sharedMaxGm[taskArgsPre.maxOutOffset]; + auto sharedSumTile = sharedSumGm[taskArgsPre.maxOutOffset]; + + epilogueSoftmax( + pL1OutTile, + qkResTensor, + expSumUb[(taskArgsPre.taskId - 1) % 3], + expSumUb[taskArgsPre.taskIdMod3], + expMaxUb[taskArgsPre.taskIdMod3], + maxUb[(taskArgsPre.taskId - 1) % 3], + maxUb[taskArgsPre.taskIdMod3], + sharedMaxTile, + sharedSumTile, + taskArgsPre.isUpdate, + taskArgsPre.isLastKv, + SharedInfer::SYNC_QK_READY_FLAG[taskArgsPre.taskIdMod2], + SharedInfer::SYNC_SOFTMAX_READY_FLAG[taskArgsPre.taskIdMod3], + SharedInfer::QK_UB_RELEASE_FLAG[taskArgsPre.taskIdMod2], + taskArgsPre.taskIdMod2, + taskArgsPre.taskIdMod3 + ); + } + } + + if (taskId > 1 && notLast) { + if ASCEND_IS_AIC { + auto &taskArgsPre2 = taskArgList[(taskId - 2) % 4]; + AscendC::CrossCoreWaitFlag(SharedInfer::SYNC_SOFTMAX_READY_FLAG[taskArgsPre2.taskIdMod3]); + AscendC::CrossCoreWaitFlag(16 + SharedInfer::SYNC_SOFTMAX_READY_FLAG[taskArgsPre2.taskIdMod3]); + + auto layoutPvRes = tla::MakeLayout(taskArgsPre2.blockQLen, embedTemplateType); + auto tensorPvRes = tla::MakeTensor(pvTensorList[taskArgsPre2.taskIdMod2], layoutPvRes, Arch::PositionUB{}); + + auto layoutPInL1 = tla::MakeLayout(qSeqlenTemplateType, kvSeqlenTemplateType); + auto tensorPInL1 = tla::MakeTensor(pL1TensorList[taskArgsPre2.taskIdMod3], layoutPInL1, Arch::PositionL1{}); + + auto tensorInV = GetTile( + tensorV, + tla::MakeCoord(taskArgsPre2.kvCoord, taskArgsPre2.kvNCoord), + tla::MakeShape(taskArgsPre2.blockKvLen, headDim) + ); + + auto actualShape = tla::MakeShape(taskArgsPre2.blockQLen, headDim, taskArgsPre2.blockKvLen); + + blockMmadPV( + tensorPInL1, tensorInV, tensorPvRes, + actualShape, taskIdL0A, taskIdL0B, SharedInfer::PV_UB_RELEASE_FLAG[taskArgsPre2.taskIdMod2] + ); + + AscendC::CrossCoreSetFlag(SharedInfer::SYNC_PV_READY_FLAG[taskArgsPre2.taskIdMod2]); + AscendC::CrossCoreSetFlag(16 + SharedInfer::SYNC_PV_READY_FLAG[taskArgsPre2.taskIdMod2]); + } + } + + if (taskId > 2) { + if ASCEND_IS_AIV { + auto &taskArgsPre3 = taskArgList[(taskId - 3) % 4]; + AscendC::CrossCoreWaitFlag(SharedInfer::SYNC_PV_READY_FLAG[taskArgsPre3.taskIdMod2]); + + auto layoutPvRes = tla::MakeLayout(taskArgsPre3.halfBlockQLen, headDim); + auto tensorPvRes = tla::MakeTensor(pvTensorList[taskArgsPre3.taskIdMod2], layoutPvRes, Arch::PositionUB{}); + + auto sharedAttnOutGmTile = GetTile( + tensorO, + tla::MakeCoord(taskArgsPre3.qCoord + taskArgsPre3.halfBlockQOffset, taskArgsPre3.qNCoord), + tla::MakeShape(taskArgsPre3.halfBlockQLen, headDim) + ); + + epilogueRescale( + sharedAttnOutGmTile, + expMaxUb[taskArgsPre3.taskIdMod3], + tensorPvRes, + taskArgsPre3.isFirstKv, + taskArgsPre3.isLastKv, + SharedInfer::PV_UB_RELEASE_FLAG[taskArgsPre3.taskIdMod2] + ); + } + } + + auto nextTaskId = (taskId + 1) % 4; + auto currentTaskId = taskId % 4; + taskArgList[nextTaskId] = taskArgList[currentTaskId]; + taskId++; + } + } + + WaitFlag(); + + // dump sharedO + // if (coreIdx == 0) { + // AscendC::printf("qHeads %d kvHeads %d headDim %d\n", qHeads, kvHeads, headDim); + // for (int i = 0; i < 8; i++) { + // AscendC::printf("token %d sharedO res\n", i); + // AscendC::DumpTensor(gSharedO[i * strideQO], 1, 8); + // } + // AscendC::printf("sharedMax res\n"); + // AscendC::DumpTensor(sharedMaxGm, 6, 8); + // AscendC::printf("sharedSum res\n"); + // AscendC::DumpTensor(sharedSumGm, 8, 8); + // } + } + + private: + static constexpr uint8_t SYNC_MODE = 4; + Arch::Resource resource; + AscendC::GlobalTensor sharedKvLensGm; + AscendC::LocalTensor qkTensorList[2]; + AscendC::LocalTensor pL1TensorList[3]; + AscendC::LocalTensor pvTensorList[2]; + AscendC::LocalTensor expSumUb[3]; + AscendC::LocalTensor expMaxUb[3]; + AscendC::LocalTensor maxUb[3]; + + SharedInfer::TaskArgs taskArgList[4]; + XAttentionTilingData* tilingData; + + int32_t batchSize{0}; + int32_t beamSize{0}; + int32_t qHeads{0}; + int32_t kvHeads{0}; + int32_t groupSize{0}; + int32_t headDim{0}; + int32_t totalTokensQ{0}; + int32_t sharedKvTokens{0}; + int64_t coreNum; + int64_t coreIdx; + int64_t subVecIdx{0}; + float scaleValue; + + uint64_t strideQO{0}; + uint64_t strideKV{0}; + uint32_t l1BufAddrStart = 0; + uint32_t l0CBufAddrStart = 0; + uint32_t ubBufAddrStart = 0; + + private: + CATLASS_DEVICE void GetQTaskInfo(SharedInfer::TaskArgs &taskArgs, int32_t qTaskId) { + int32_t perBatchHeadTaskNum = tilingData->sharedInfo.perBatchHeadTaskNum; + int32_t qBlockId = qTaskId % perBatchHeadTaskNum; + int32_t outerId = qTaskId / perBatchHeadTaskNum; + int32_t qHeadId = outerId % qHeads; + int32_t batchId = outerId / qHeads; + taskArgs.batchId = batchId; + taskArgs.qHeadId = qHeadId; + taskArgs.kvHeadId = qHeadId / groupSize; + taskArgs.qBlockId = qBlockId; + taskArgs.blockQLen = qBlockId == (perBatchHeadTaskNum - 1) ? (beamSize - qBlockId * qSeqlenTemplateType) : qSeqlenTemplateType; + taskArgs.actualKvLen = sharedKvLensGm.GetValue(batchId); + taskArgs.qCoord = batchId * beamSize + qBlockId * qSeqlenTemplateType; + taskArgs.qNCoord = qHeadId * headDim; + taskArgs.kvNCoord = taskArgs.kvHeadId * headDim; + + if ASCEND_IS_AIV { + int32_t halfQLen = (taskArgs.blockQLen + 1) / 2; + taskArgs.halfBlockQLen = (subVecIdx == 0) ? halfQLen : (taskArgs.blockQLen - halfQLen); + taskArgs.halfBlockQOffset = (subVecIdx == 0) ? 0 : halfQLen; + taskArgs.maxOutOffset = (taskArgs.qCoord + taskArgs.halfBlockQOffset) * qHeads + qHeadId; + } + + int32_t batchOffset = 0; + for (int bId = 0; bId < batchId; bId++) { + batchOffset += sharedKvLensGm.GetValue(bId); + } + + taskArgs.kvBatchOffset = batchOffset; + } + + CATLASS_DEVICE void GetKvTaskInfo(SharedInfer::TaskArgs &taskArgs, int32_t kvBlockId, int32_t kvBlockNum, int32_t taskId) { + auto actualKvLen = taskArgs.actualKvLen; + bool isFirstKv = kvBlockId == 0; + bool isUpdate = kvBlockId > 0; + bool isLastKv = kvBlockId == kvBlockNum - 1; + taskArgs.taskId = taskId; + taskArgs.kvBlockId = kvBlockId; + taskArgs.blockKvLen = isLastKv ? (actualKvLen - kvBlockId * kvSeqlenTemplateType) : kvSeqlenTemplateType; + taskArgs.kvCoord = taskArgs.kvBatchOffset + kvBlockId * kvSeqlenTemplateType; + taskArgs.isFirstKv = isFirstKv; + taskArgs.isUpdate = isUpdate; + taskArgs.isLastKv = isLastKv; + taskArgs.taskIdMod2 = taskId % 2; + taskArgs.taskIdMod3 = taskId % 3; + } + + CATLASS_DEVICE void SetFlag() { + if ASCEND_IS_AIC { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + } else { + AscendC::CrossCoreSetFlag(SharedInfer::QK_UB_RELEASE_FLAG[0]); + AscendC::CrossCoreSetFlag(SharedInfer::QK_UB_RELEASE_FLAG[1]); + AscendC::CrossCoreSetFlag(SharedInfer::PV_UB_RELEASE_FLAG[0]); + AscendC::CrossCoreSetFlag(SharedInfer::PV_UB_RELEASE_FLAG[1]); + } + } + + CATLASS_DEVICE void WaitFlag() { + if ASCEND_IS_AIC { + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + AscendC::CrossCoreWaitFlag(SharedInfer::QK_UB_RELEASE_FLAG[0]); + AscendC::CrossCoreWaitFlag(16 + SharedInfer::QK_UB_RELEASE_FLAG[0]); + AscendC::CrossCoreWaitFlag(SharedInfer::QK_UB_RELEASE_FLAG[1]); + AscendC::CrossCoreWaitFlag(16 + SharedInfer::QK_UB_RELEASE_FLAG[1]); + AscendC::CrossCoreWaitFlag(SharedInfer::PV_UB_RELEASE_FLAG[0]); + AscendC::CrossCoreWaitFlag(16 + SharedInfer::PV_UB_RELEASE_FLAG[0]); + AscendC::CrossCoreWaitFlag(SharedInfer::PV_UB_RELEASE_FLAG[1]); + AscendC::CrossCoreWaitFlag(16 + SharedInfer::PV_UB_RELEASE_FLAG[1]); + } + } +}; + +#endif \ No newline at end of file diff --git a/xllm_ops/x_attention/op_kernel/arch35/unshared_infer_catlass_kernel.h b/xllm_ops/x_attention/op_kernel/arch35/unshared_infer_catlass_kernel.h new file mode 100644 index 0000000..915ca86 --- /dev/null +++ b/xllm_ops/x_attention/op_kernel/arch35/unshared_infer_catlass_kernel.h @@ -0,0 +1,375 @@ + +#ifndef X_ATTN_UNSHARED_FA_INFER_CATLASS_KERNEL_H +#define X_ATTN_UNSHARED_FA_INFER_CATLASS_KERNEL_H + +#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 "catlass/epilogue/dispatch_policy.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" +#include "x_attention_common.h" +#include "kernel_operator.h" + +using namespace Catlass; + +template < + class BlockMmadQK, + class BlockMmadPV, + class EpiloueSoftmax, + typename KVLEN_T, + typename TABLE_T> +class UnSharedInferKernel { + public: + using ArchTag = typename BlockMmadQK::ArchTag; + using L1TileShape = typename BlockMmadQK::L1TileShape; + using ElementQ = typename BlockMmadQK::ElementA; + using LayoutTagQ = typename BlockMmadQK::LayoutTagA; + using ElementK = typename BlockMmadQK::ElementB; + using LayoutTagK = typename BlockMmadQK::LayoutTagB; + using ElementS = typename BlockMmadQK::ElementC; + using LayoutTagS = typename BlockMmadQK::LayoutTagC; + + using ElementP = typename BlockMmadPV::ElementA; + using LayoutTagP = typename BlockMmadPV::LayoutTagA; + using LayoutTagPL1 = typename BlockMmadPV::TileCopy::LayoutTagL1A; + using ElementV = typename BlockMmadPV::ElementB; + using LayoutTagV = typename BlockMmadPV::LayoutTagB; + using ElementOTmp = typename BlockMmadPV::ElementC; + using LayoutTagOTmp = typename BlockMmadPV::LayoutTagC; + + static constexpr uint32_t qSeqlenTemplateType = tla::get<0>(L1TileShape{}); + static constexpr uint32_t kvSeqlenTemplateType = tla::get<1>(L1TileShape{}); + static constexpr uint32_t embedTemplateType = tla::get<2>(L1TileShape{}); + static constexpr uint32_t halfQSeqlenTemplateType = qSeqlenTemplateType / CV_RATIO; + + CATLASS_DEVICE + UnSharedInferKernel(XAttentionTilingData *tilingData) { + this->tilingData = tilingData; + } + + template + CATLASS_DEVICE void operator()(XAttnKernelCommonParams const ¶ms); + + CATLASS_DEVICE void Init(XAttnKernelCommonParams const ¶ms) { + auto qkSize = halfQSeqlenTemplateType * kvSeqlenTemplateType * sizeof(ElementS); + // AscendC::printf("qkSize %d ubBufAddrStart %d\n", qkSize, ubBufAddrStart); + + for (int i = 0; i < 2; i++) { + qkTensorList[i] = resource.ubBuf.template GetBufferByByte(ubBufAddrStart); + ubBufAddrStart += qkSize; + } + auto pL1Size = qSeqlenTemplateType * kvSeqlenTemplateType * sizeof(ElementP); + for (int i = 0; i < 3; i++) { + pL1TensorList[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); + l1BufAddrStart += pL1Size; + } + + decodeStepGm.SetGlobalBuffer((__gm__ KVLEN_T *)params.decodeStep); + unsharedKvSeqLen = static_cast(decodeStepGm.GetValue(0)); + blockTableGm.SetGlobalBuffer((__gm__ TABLE_T *)params.unsharedBlockTable); + + batchSize = tilingData->baseInfo.batchSize; + beamSize = tilingData->baseInfo.beamSize; + qHeads = tilingData->baseInfo.qHeads; + kvHeads = tilingData->baseInfo.kvHeads; + groupSize = tilingData->baseInfo.groupSize; + headDim = tilingData->baseInfo.headDim; + scaleValue = tilingData->baseInfo.scaleValue; + totalTokensQ = tilingData->baseInfo.totalTokensQ; + maxDecodeStep = tilingData->baseInfo.maxDecodeStep; + + kvBatchStride = tilingData->unsharedInfo.kvBatchStride; + groupCountPerLoop = tilingData->unsharedInfo.groupCountPerLoop; + perBatchTaskNum = tilingData->unsharedInfo.perBatchTaskNum; + perCoreTaskNum = tilingData->unsharedInfo.perCoreTaskNum; + totalTaskNum = tilingData->unsharedInfo.totalTaskNum; + coreNum = tilingData->unsharedInfo.usedCoreNum; + + coreIdx = AscendC::GetBlockIdx(); + subVecIdx = AscendC::GetSubBlockIdx(); + if ASCEND_IS_AIV { + coreIdx = coreIdx / CV_RATIO; + int32_t halfGroupCount = (groupCountPerLoop + CV_RATIO - 1) / CV_RATIO; + halfVecGroupCount = (subVecIdx == 0) ? halfGroupCount : (groupCountPerLoop - halfGroupCount); + halfVecGroupOffset = (subVecIdx == 0) ? 0 : halfGroupCount; + halfVecRowCount = halfVecGroupCount * groupSize; + halfVecRowOffset = halfVecGroupOffset * groupSize; + } + coreIdx = coreIdx - tilingData->sharedInfo.usedCoreNum; + blockQLen = groupCountPerLoop * groupSize; + blockKvLen = groupCountPerLoop * maxDecodeStep; + } + + CATLASS_DEVICE void operator()(XAttnKernelCommonParams const ¶ms) { + uint32_t taskIdL0A = 0; + uint32_t taskIdL0B = 0; + uint32_t taskIdL0C = 0; + + Init(params); + SetFlag(); + + AscendC::GlobalTensor gQ; + gQ.SetGlobalBuffer((__gm__ ElementQ *)params.q); + auto layoutQ = tla::MakeLayout(totalTokensQ * qHeads, headDim); + auto tensorQ = tla::MakeTensor(gQ, layoutQ, Arch::PositionGM{}); + AscendC::GlobalTensor gK; + gK.SetGlobalBuffer((__gm__ ElementK *)params.unsharedK); + AscendC::GlobalTensor gV; + gV.SetGlobalBuffer((__gm__ ElementV *)params.unsharedV); + + AscendC::GlobalTensor gUnSharedO; + gUnSharedO.SetGlobalBuffer((__gm__ ElementOTmp *)params.unsharedO); + auto layoutO = tla::MakeLayout(totalTokensQ * qHeads, headDim); + auto tensorO = tla::MakeTensor(gUnSharedO, layoutO, Arch::PositionGM{}); + + AscendC::GlobalTensor unsharedMaxGm; + unsharedMaxGm.SetGlobalBuffer((__gm__ ElementOTmp *)params.unsharedMax); + AscendC::GlobalTensor unsharedSumGm; + unsharedSumGm.SetGlobalBuffer((__gm__ ElementOTmp *)params.unsharedSum); + + BlockMmadQK blockMmadQK(resource, l1BufAddrStart, l0CBufAddrStart); + BlockMmadPV blockMmadPV(resource, l1BufAddrStart, l0CBufAddrStart); + EpiloueSoftmax epilogueSoftmax(resource, ubBufAddrStart, scaleValue, unsharedKvSeqLen, maxDecodeStep, groupCountPerLoop, groupSize); + + int32_t batchKvLen = beamSize * kvHeads * maxDecodeStep; + int32_t taskId = 0; + int32_t perCoreTaskNum = tilingData->unsharedInfo.perCoreTaskNum; + int32_t totalTaskNum = tilingData->unsharedInfo.totalTaskNum; + + int32_t taskStartId = coreIdx * perCoreTaskNum; + if (taskStartId >= totalTaskNum) { + WaitFlag(); + return; + } + int32_t coreTaskNum = perCoreTaskNum; + int32_t taskEndId = taskStartId + coreTaskNum; + + if (taskEndId > totalTaskNum) { + taskEndId = totalTaskNum; + coreTaskNum = taskEndId - taskStartId; + } + + // AscendC::printf("coreIdx %d taskStartId %d taskEndId %d coreTaskNum %d\n", coreIdx, taskStartId, taskEndId, coreTaskNum); + + for (int32_t groupTaskId = taskStartId; groupTaskId < taskEndId + 2; groupTaskId++) + { + bool notLastTwoLoop = groupTaskId < taskEndId; + bool notLast = groupTaskId < taskEndId + 1; + + if (notLastTwoLoop) { + UnSharedInfer::TaskArgs taskArgs; + GetTaskInfo(taskArgs, groupTaskId, taskId); + taskArgList[taskId % 3] = taskArgs; + if ASCEND_IS_AIC { + auto nowTaskId = taskId % 3; + UnSharedInfer::TaskArgs &taskArgsNow = taskArgList[nowTaskId]; + auto actualShape = tla::MakeShape(blockQLen, blockKvLen, headDim); + auto tensorQTile = GetTile( + tensorQ, + tla::MakeCoord(taskArgsNow.qCoord, 0), + tla::MakeShape(blockQLen, headDim) + ); + + auto layoutK = tla::MakeLayout(headDim, batchKvLen); + auto tensorK = tla::MakeTensor(gK[taskArgsNow.cacheBlockId * kvBatchStride], layoutK, Arch::PositionGM{}); + auto tensorKTile = GetTile( + tensorK, + tla::MakeCoord(0, taskArgsNow.kvCoord), + tla::MakeShape(headDim, blockKvLen) + ); + + auto layoutQKRes = tla::MakeLayout(blockQLen, kvSeqlenTemplateType); + auto tensorQKRes = tla::MakeTensor(qkTensorList[taskArgsNow.taskIdMod2], layoutQKRes, Arch::PositionUB{}); + + blockMmadQK( + tensorQTile, tensorKTile, tensorQKRes, actualShape, + UnSharedInfer::QK_UB_RELEASE_FLAG[taskArgsNow.taskIdMod2], + taskIdL0A, taskIdL0B, taskIdL0C + ); + AscendC::CrossCoreSetFlag(UnSharedInfer::SYNC_QK_READY_FLAG[taskArgsNow.taskIdMod2]); + AscendC::CrossCoreSetFlag(16 + UnSharedInfer::SYNC_QK_READY_FLAG[taskArgsNow.taskIdMod2]); + + } + } + + if (taskId > 0 && notLast) { + if ASCEND_IS_AIV { + auto &taskArgsPre = taskArgList[(taskId - 1) % 3]; + auto qkResLayout = tla::MakeLayout(halfVecRowCount, blockKvLen); + auto qkResTensor = tla::MakeTensor(qkTensorList[taskArgsPre.taskIdMod2], qkResLayout, Arch::PositionUB{}); + auto pL1OutLayout = tla::MakeLayout(qSeqlenTemplateType, kvSeqlenTemplateType); + auto pL1OutTensor = tla::MakeTensor(pL1TensorList[taskArgsPre.taskIdMod3], pL1OutLayout, Arch::PositionL1{}); + auto pL1OutTile = GetTile( + pL1OutTensor, + tla::MakeCoord(halfVecRowOffset, 0), + tla::MakeShape(halfVecRowCount, kvSeqlenTemplateType) + ); + auto unsharedMaxTile = unsharedMaxGm[taskArgsPre.maxOutOffset]; + auto unsharedSumTile = unsharedSumGm[taskArgsPre.maxOutOffset]; + + epilogueSoftmax( + pL1OutTile, + qkResTensor, + unsharedMaxTile, + unsharedSumTile, + UnSharedInfer::SYNC_QK_READY_FLAG[taskArgsPre.taskIdMod2], + UnSharedInfer::SYNC_SOFTMAX_READY_FLAG[taskArgsPre.taskIdMod3], + UnSharedInfer::QK_UB_RELEASE_FLAG[taskArgsPre.taskIdMod2], + taskArgsPre.taskIdMod2, + taskArgsPre.taskIdMod3 + ); + + } + } + + if (taskId > 1) { + if ASCEND_IS_AIC { + auto &taskArgsPre2 = taskArgList[(taskId - 2) % 3]; + AscendC::CrossCoreWaitFlag(UnSharedInfer::SYNC_SOFTMAX_READY_FLAG[taskArgsPre2.taskIdMod3]); + AscendC::CrossCoreWaitFlag(16 + UnSharedInfer::SYNC_SOFTMAX_READY_FLAG[taskArgsPre2.taskIdMod3]); + + auto tensorOTile = GetTile( + tensorO, + tla::MakeCoord(taskArgsPre2.qCoord, 0), + tla::MakeShape(blockQLen, headDim) + ); + + auto layoutPInL1 = tla::MakeLayout(qSeqlenTemplateType, kvSeqlenTemplateType); + auto tensorPInL1 = tla::MakeTensor(pL1TensorList[taskArgsPre2.taskIdMod3], layoutPInL1, Arch::PositionL1{}); + + auto layoutV = tla::MakeLayout(batchKvLen, headDim); + auto tensorV = tla::MakeTensor(gV[taskArgsPre2.cacheBlockId * kvBatchStride], layoutV, Arch::PositionGM{}); + auto tensorVTile = GetTile( + tensorV, + tla::MakeCoord(taskArgsPre2.kvCoord, 0), + tla::MakeShape(blockKvLen, headDim) + ); + + auto actualShape = tla::MakeShape(blockQLen, headDim, blockKvLen); + + blockMmadPV( + tensorPInL1, tensorVTile, tensorOTile, + actualShape, taskIdL0A, taskIdL0B, taskIdL0C + ); + } + } + + auto nextTaskId = (taskId + 1) % 3; + auto currentTaskId = taskId % 3; + taskArgList[nextTaskId] = taskArgList[currentTaskId]; + taskId++; + } + // if (coreIdx == 0) { + // AscendC::printf("qHeads %d kvHeads %d headDim %d\n", qHeads, kvHeads, headDim); + // for (int i = 64; i < 80; i++) { + // AscendC::printf("token %d unsharedO res\n", i); + // AscendC::DumpTensor(gUnSharedO[i * headDim], 1, 8); + // } + // AscendC::printf("unsharedMax res\n"); + // AscendC::DumpTensor(unsharedMaxGm, 6, 8); + // AscendC::printf("unsharedSum res\n"); + // AscendC::DumpTensor(unsharedSumGm, 8, 8); + // } + + WaitFlag(); + } + + private: + static constexpr uint8_t SYNC_MODE = 4; + Arch::Resource resource; + AscendC::GlobalTensor decodeStepGm; + AscendC::LocalTensor qkTensorList[2]; + AscendC::LocalTensor pL1TensorList[3]; + AscendC::GlobalTensor blockTableGm; + + UnSharedInfer::TaskArgs taskArgList[3]; + XAttentionTilingData* tilingData; + + int32_t batchSize{0}; + int32_t beamSize{0}; + int32_t qHeads{0}; + int32_t kvHeads{0}; + int32_t groupSize{0}; + int32_t headDim{0}; + int32_t totalTokensQ{0}; + int32_t unsharedKvSeqLen{0}; + int32_t maxDecodeStep{0}; + int32_t groupCountPerLoop{0}; + int32_t kvBatchStride; + int32_t perBatchTaskNum{0}; + int32_t perCoreTaskNum{0}; + int32_t totalTaskNum{0}; + int32_t halfVecGroupCount{0}; + int32_t halfVecGroupOffset{0}; + int32_t halfVecRowCount; + int32_t halfVecRowOffset; + int32_t blockQLen; + int32_t blockKvLen; + int32_t coreNum; + int64_t coreIdx; + int64_t subVecIdx{0}; + float scaleValue; + + uint32_t l1BufAddrStart = 0; + uint32_t l0CBufAddrStart = 0; + uint32_t ubBufAddrStart = 0; + + private: + CATLASS_DEVICE void GetTaskInfo(UnSharedInfer::TaskArgs &taskArgs, int32_t groupTaskId, int32_t taskId) { + taskArgs.taskId = taskId; + taskArgs.taskIdMod2 = taskId % 2; + taskArgs.taskIdMod3 = taskId % 3; + + int32_t batchId = groupTaskId / perBatchTaskNum; + int32_t cacheBlockId = blockTableGm.GetValue(batchId); + int32_t groupCountBlockId = groupTaskId % perBatchTaskNum; + + taskArgs.batchId = batchId; + taskArgs.cacheBlockId = cacheBlockId; + taskArgs.groupCountBlockId = groupCountBlockId; + taskArgs.qCoord = batchId * beamSize * qHeads + groupCountBlockId * groupCountPerLoop * groupSize; + taskArgs.kvCoord = groupCountBlockId * groupCountPerLoop * maxDecodeStep; + + if ASCEND_IS_AIV { + taskArgs.maxOutOffset = taskArgs.qCoord + halfVecRowOffset; + } + + } + + CATLASS_DEVICE void SetFlag() { + if ASCEND_IS_AIC { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + } else { + // AscendC::printf("coreIdx %d subVecIdx %d set qk_ub_flag %d %d \n", coreIdx, subVecIdx, UnSharedInfer::QK_UB_RELEASE_FLAG[0], UnSharedInfer::QK_UB_RELEASE_FLAG[1]); + AscendC::CrossCoreSetFlag(UnSharedInfer::QK_UB_RELEASE_FLAG[0]); + AscendC::CrossCoreSetFlag(UnSharedInfer::QK_UB_RELEASE_FLAG[1]); + } + } + + CATLASS_DEVICE void WaitFlag() { + if ASCEND_IS_AIC { + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::CrossCoreWaitFlag(UnSharedInfer::QK_UB_RELEASE_FLAG[0]); + AscendC::CrossCoreWaitFlag(UnSharedInfer::QK_UB_RELEASE_FLAG[1]); + AscendC::CrossCoreWaitFlag(16 + UnSharedInfer::QK_UB_RELEASE_FLAG[0]); + AscendC::CrossCoreWaitFlag(16 + UnSharedInfer::QK_UB_RELEASE_FLAG[1]); + } + } +}; + +#endif \ No newline at end of file diff --git a/xllm_ops/x_attention/op_kernel/arch35/x_attention_catlass_helper.h b/xllm_ops/x_attention/op_kernel/arch35/x_attention_catlass_helper.h new file mode 100644 index 0000000..05e4e75 --- /dev/null +++ b/xllm_ops/x_attention/op_kernel/arch35/x_attention_catlass_helper.h @@ -0,0 +1,120 @@ + + +#ifndef X_ATTN_CATLASS_HELPER_H +#define X_ATTN_CATLASS_HELPER_H +#include "shared_infer_catlass_kernel.h" +#include "unshared_infer_catlass_kernel.h" +#include "combine_kernel.h" + +template +CATLASS_DEVICE void CallSharedInferKernel(const XAttnKernelCommonParams& params, XAttentionTilingData* tilingData) { + using ArchTag = Arch::Ascend950; + using ElementQ = INPUT_T; + using LayoutQ = layout::RowMajor; + using ElementK = INPUT_T; + using LayoutK = layout::ColumnMajor; + using ElementV = INPUT_T; + using LayoutV = layout::RowMajor; + using ElementS = float; + using LayoutS = layout::RowMajor; + using ElementP = INPUT_T; + using LayoutP = layout::RowMajor; + using ElementOTmp = float; + using LayoutOTmp = layout::RowMajor; + // L1TileShape::K must be embdding + using L1TileShape = tla::Shape<_128, _128, _128>; + using L0TileShape = L1TileShape; + // GEMM Block, implement Q @ K^T of Flash Attention Infer + using DispatchPolicyQK = Gemm::MmadXASharedQK; + using TileCopyQK = Gemm::Tile::PackedTileCopyTlaToUB< + ArchTag, ElementQ, LayoutQ, ElementK, LayoutK, ElementS, LayoutS, void, Gemm::Tile::CopyL0CToUBMode::SPLIT_M>; + using TileMmadQK = Gemm::Tile::TileMmadTla; + using BlockMmadQK = Gemm::Block::BlockMmadTla; + + // Shared Epilogue Block, update rowsum rowmax and copyOut on lastStackTile + using DispatchPolicyOnlineSoftmax = Epilogue::EpilogueAscend950XASharedSoftmax; + using PType = Gemm::GemmType; + using SType = Gemm::GemmType; + using EpilogueOnlineSoftmax = Epilogue::Block::BlockEpilogue; + + // GEMM Block, implement P @ V of Flash Attention Infer + using DispatchPolicyPV = Gemm::MmadXASharedPV; + using TileCopyPV = Gemm::Tile::PackedTileCopyTlaToUB< + ArchTag, ElementP, LayoutP, ElementV, LayoutV, ElementOTmp, LayoutOTmp, void, Gemm::Tile::CopyL0CToUBMode::SPLIT_M>; + using TileMmadPV = Gemm::Tile::TileMmadTla; + using BlockMmadPV = Gemm::Block::BlockMmadTla; + + // Shared Epilogue RescaleO，do not div rowSum or cast on lastStackTile + using DispatchPolicyRescaleO = Epilogue::EpilogueAscend950XASharedRescaleO; + using OTmpType = Gemm::GemmType; + using EpilogueRescaleO = Epilogue::Block::BlockEpilogue; + + using SharedFAInferKernel = SharedFaInferKernel< + BlockMmadQK, BlockMmadPV, EpilogueOnlineSoftmax, EpilogueRescaleO, KVLEN_T>; + + SharedFAInferKernel sharedInferKernel(tilingData); + sharedInferKernel(params); +} + +template +CATLASS_DEVICE void CallUnsharedInferKernel(const XAttnKernelCommonParams& params, XAttentionTilingData* tilingData) { + using ArchTag = Arch::Ascend950; + using ElementQ = INPUT_T; + using LayoutQ = layout::RowMajor; + using ElementK = INPUT_T; + using LayoutK = layout::ColumnMajor; + using ElementV = INPUT_T; + using LayoutV = layout::RowMajor; + using ElementS = float; + using LayoutS = layout::RowMajor; + using ElementP = INPUT_T; + using LayoutP = layout::RowMajor; + using ElementOTmp = float; + using LayoutOTmp = layout::RowMajor; + // L1TileShape::K must be embdding + using L1TileShape = tla::Shape<_128, _128, _128>; + using L0TileShape = L1TileShape; + // GEMM Block, implement Q @ K^T of Flash Attention Infer + using DispatchPolicyQK = Gemm::MmadXAUnsharedQK; + using TileCopyQK = Gemm::Tile::PackedTileCopyTlaToUB< + ArchTag, ElementQ, LayoutQ, ElementK, LayoutK, ElementS, LayoutS, void, Gemm::Tile::CopyL0CToUBMode::SPLIT_M>; + using TileMmadQK = Gemm::Tile::TileMmadTla; + using BlockMmadQK = Gemm::Block::BlockMmadTla; + + // Shared Epilogue Block, update rowsum rowmax and copyOut on lastStackTile + using DispatchPolicySoftmax = Epilogue::EpilogueAscend950XAUnsharedSoftmax; + using PType = Gemm::GemmType; + using SType = Gemm::GemmType; + using EpilogueSoftmax = Epilogue::Block::BlockEpilogue; + + // GEMM Block, implement P @ V of Flash Attention Infer + using DispatchPolicyPV = Gemm::MmadXAUnsharedPV; + using TileCopyPV = Gemm::Tile::PackedTileCopyTla< + ArchTag, ElementP, LayoutP, ElementV, LayoutV, ElementOTmp, LayoutOTmp>; + using TileMmadPV = Gemm::Tile::TileMmadTla; + using BlockMmadPV = Gemm::Block::BlockMmadTla; + + using UnSharedInferKernel = UnSharedInferKernel< + BlockMmadQK, BlockMmadPV, EpilogueSoftmax, KVLEN_T, TABLE_T>; + + UnSharedInferKernel unsharedInferKernel(tilingData); + unsharedInferKernel(params); +} + + +template +CATLASS_DEVICE void CallCombineScale(const XAttnKernelCommonParams& params, XAttentionTilingData* tilingData) { + using DispatchPolicyCombine = Epilogue::EpilogueAscend950XACombineScale; + using ElementInput = float; + using LayoutInput = layout::RowMajor; + using ElementOutput = INPUT_T; + using LayoutOutput = layout::RowMajor; + using InputType = Gemm::GemmType; + using OutputType = Gemm::GemmType; + using EpilogueCombineScale = Epilogue::Block::BlockEpilogue; + + using CombineKernel = CombineScaleKernel; + CombineKernel combineKernel(tilingData); + combineKernel(params); +} +#endif diff --git a/xllm_ops/x_attention/op_kernel/arch35/x_attention_common.h b/xllm_ops/x_attention/op_kernel/arch35/x_attention_common.h new file mode 100644 index 0000000..35b58d4 --- /dev/null +++ b/xllm_ops/x_attention/op_kernel/arch35/x_attention_common.h @@ -0,0 +1,109 @@ + +#ifndef X_ATTENTION_COMMON +#define X_ATTENTION_COMMON + + +constexpr uint32_t BLOCK_SIZE = 16; +constexpr uint32_t CV_RATIO = 2; + +namespace SharedInfer { + constexpr uint16_t SYNC_QK_READY_FLAG[2] = {0, 1}; + constexpr uint16_t SYNC_SOFTMAX_READY_FLAG[3] = {2, 3, 4}; + constexpr uint16_t SYNC_PV_READY_FLAG[2] = {5, 6}; + constexpr uint16_t QK_UB_RELEASE_FLAG[2] = {7, 8}; + constexpr uint16_t PV_UB_RELEASE_FLAG[2] = {9, 10}; + constexpr uint32_t COMPUTE_PIPE_NUM = 3; + struct TaskArgs { + int32_t taskId = 0; + int32_t batchId = 0; + int32_t qHeadId = 0; + int32_t kvHeadId = 0; + int32_t qBlockId = 0; + int32_t kvBlockId = 0; + int32_t actualKvLen = 0; + int32_t blockQLen = 0; + int32_t blockKvLen = 0; + int32_t qCoord = 0; + int32_t kvCoord = 0; + int32_t qNCoord = 0; + int32_t kvNCoord = 0; + bool isFirstKv = false; + bool isUpdate = false; + bool isLastKv = false; + int32_t taskIdMod2 = 0; + int32_t taskIdMod3 = 0; + int32_t kvBatchOffset = 0; + int32_t halfBlockQLen = 0; + int32_t halfBlockQOffset = 0; + int32_t maxOutOffset = 0; + }; +} + +namespace UnSharedInfer { + constexpr uint16_t SYNC_QK_READY_FLAG[2] = {0, 1}; + constexpr uint16_t SYNC_SOFTMAX_READY_FLAG[3] = {2, 3, 4}; + constexpr uint16_t QK_UB_RELEASE_FLAG[2] = {5, 6}; + constexpr uint32_t COMPUTE_PIPE_NUM = 3; + struct TaskArgs { + int32_t taskId; + int32_t batchId; + int32_t cacheBlockId; + int32_t groupCountBlockId; + int32_t qCoord; + int32_t kvCoord; + int32_t taskIdMod2; + int32_t taskIdMod3; + int32_t maxOutOffset = 0; + }; +} + +struct XAttnKernelCommonParams { + GM_ADDR q; + GM_ADDR sharedK; + GM_ADDR sharedV; + GM_ADDR unsharedK; + GM_ADDR unsharedV; + GM_ADDR sharedBlockTable; + GM_ADDR unsharedBlockTable; + GM_ADDR sharedKvLens; // shared Kv + GM_ADDR decodeStep; // unshared kv: 1, 2, 3 + GM_ADDR sharedO; + GM_ADDR sharedMax; + GM_ADDR sharedSum; + GM_ADDR unsharedO; + GM_ADDR unsharedMax; + GM_ADDR unsharedSum; + GM_ADDR o; // final combine out + GM_ADDR tiling; + + CATLASS_DEVICE + XAttnKernelCommonParams() { + } + + CATLASS_DEVICE + XAttnKernelCommonParams( + GM_ADDR q_, GM_ADDR sharedK_, GM_ADDR sharedV_, GM_ADDR unsharedK_, GM_ADDR unsharedV_, + GM_ADDR sharedBlockTable_, GM_ADDR unsharedBlockTable_, GM_ADDR sharedKvLens_, GM_ADDR decodeStep_, + GM_ADDR sharedO_, GM_ADDR sharedMax_, GM_ADDR sharedSum_, GM_ADDR unsharedO_, GM_ADDR unsharedMax_, + GM_ADDR unsharedSum_, GM_ADDR o_, GM_ADDR tiling_) + : q(q_), + sharedK(sharedK_), + sharedV(sharedV_), + unsharedK(unsharedK_), + unsharedV(unsharedV_), + sharedBlockTable(sharedBlockTable_), + unsharedBlockTable(unsharedBlockTable_), + sharedKvLens(sharedKvLens_), + decodeStep(decodeStep_), + sharedO(sharedO_), + sharedMax(sharedMax_), + sharedSum(sharedSum_), + unsharedO(unsharedO_), + unsharedMax(unsharedMax_), + unsharedSum(unsharedSum_), + o(o_), + tiling(tiling_) + {} +}; + +#endif diff --git a/xllm_ops/x_attention/op_kernel/x_attention.cpp b/xllm_ops/x_attention/op_kernel/x_attention.cpp index d9d4bc8..8ceabd0 100644 --- a/xllm_ops/x_attention/op_kernel/x_attention.cpp +++ b/xllm_ops/x_attention/op_kernel/x_attention.cpp @@ -13,27 +13,80 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +// A5(Ascend950/DAV_3510) arch guard. +// Device side must use __NPU_ARCH__ (per catlass migration guide); host side uses +// CATLASS_ARCH. Accept either so the A5 path is selected regardless of which macro +// the toolchain injects for the kernel translation unit. +#if (defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)) || (defined(CATLASS_ARCH) && (CATLASS_ARCH == 3510)) +#define XA_ARCH35 1 +#endif + +// Device(kernel) side lacks -DCATLASS_ARCH (host-only inject). Derive it from +// __NPU_ARCH__ HERE, before ANY include, so every catlass forwarding header in +// this translation unit dispatches to the ascend950 specialization consistently. +#if defined(XA_ARCH35) && !defined(CATLASS_ARCH) +#define CATLASS_ARCH 3510 +#endif + +// A3(AtlasA2/A3, __NPU_ARCH__ == 2201) arch guard. Derive CATLASS_ARCH so the A3 +// device translation unit resolves the catlass forwarding headers consistently. +#if !defined(XA_ARCH35) && !defined(CATLASS_ARCH) && defined(__NPU_ARCH__) && (__NPU_ARCH__ == 2201) +#define CATLASS_ARCH 2201 +#endif + #include "kernel_operator.h" -#include "x_attention_catlass_helper.h" #include "lib/matmul_intf.h" -#define CALL_XATTN_KERNEL(INPUT_TYPE, SHARED_PAGED_FLAG, UNSHARED_PAGED_FLAG) \ - do { \ - if (coreIdx < tiling_data.sharedCoreNum) { \ - CallSharedInferKernelShort(params, &tiling_data); \ - } else { \ - CallUnsharedInferKernel(params, &tiling_data); \ - } \ - AscendC::SyncAll(); \ - CallCombineScale(params, &tiling_data); \ - } while (0) +#if defined(XA_ARCH35) +#include "arch35/x_attention_catlass_helper.h" +#else +#include "x_attention_catlass_helper.h" +#endif using namespace AscendC; -extern "C" __global__ __aicore__ void x_attention(GM_ADDR query, GM_ADDR shared_key_block, GM_ADDR shared_value_block, +extern "C" __global__ __aicore__ void x_attention(GM_ADDR query, GM_ADDR shared_key_block, GM_ADDR shared_value_block, GM_ADDR unshared_key_block, GM_ADDR unshared_value_block, GM_ADDR unshared_block_table, GM_ADDR shared_kv_lens, GM_ADDR decode_step, GM_ADDR shared_block_table, GM_ADDR attn_out, GM_ADDR workspace, GM_ADDR tiling) { +#if defined(XA_ARCH35) + // ===== A5(Ascend950/DAV_3510) path ===== + // workspace layout: [sharedO, sharedMax, sharedSum, unsharedO, unsharedMax, unsharedSum] + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + GET_TILING_DATA(tiling_data, tiling); + + GM_ADDR sharedO = workspace; + GM_ADDR sharedMax = sharedO + tiling_data.qOSize; + GM_ADDR sharedSum = sharedMax + tiling_data.sumMaxSize; + GM_ADDR unsharedO = sharedSum + tiling_data.sumMaxSize; + GM_ADDR unsharedMax = unsharedO + tiling_data.qOSize; + GM_ADDR unsharedSum = unsharedMax + tiling_data.sumMaxSize; + int64_t coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(); + + XAttnKernelCommonParams params{query, shared_key_block, shared_value_block, unshared_key_block, unshared_value_block, + shared_block_table, unshared_block_table, shared_kv_lens, decode_step, sharedO, sharedMax, sharedSum, unsharedO, unsharedMax, + unsharedSum, attn_out, tiling}; + + if (coreIdx < tiling_data.sharedInfo.usedCoreNum) { + CallSharedInferKernel(params, &tiling_data); + } else { + CallUnsharedInferKernel(params, &tiling_data); + } + AscendC::SyncAll(); + CallCombineScale(params, &tiling_data); +#else + // ===== A3(AtlasA2/A3) path ===== // workspace use; [s,p,oTemp,oUpdate,shared_workspace,unshared_workspace] + #define CALL_XATTN_KERNEL(INPUT_TYPE, SHARED_PAGED_FLAG, UNSHARED_PAGED_FLAG) \ + do { \ + if (coreIdx < tiling_data.sharedCoreNum) { \ + CallSharedInferKernelShort(params, &tiling_data); \ + } else { \ + CallUnsharedInferKernel(params, &tiling_data); \ + } \ + AscendC::SyncAll(); \ + CallCombineScale(params, &tiling_data); \ + } while (0) + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); GET_TILING_DATA(tiling_data, tiling); @@ -45,8 +98,8 @@ extern "C" __global__ __aicore__ void x_attention(GM_ADDR query, GM_ADDR shared_ GM_ADDR unshared_workspace = shared_workspace + tiling_data.sharedWorkspaceSize; int64_t coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(); - XAttnKernelParams params{query, shared_key_block, shared_value_block, unshared_key_block, unshared_value_block, - shared_block_table, unshared_block_table, shared_kv_lens, decode_step, s, p, oTemp, oUpdate, shared_workspace, + XAttnKernelParams params{query, shared_key_block, shared_value_block, unshared_key_block, unshared_value_block, + shared_block_table, unshared_block_table, shared_kv_lens, decode_step, s, p, oTemp, oUpdate, shared_workspace, unshared_workspace, attn_out, tiling}; if (TILING_KEY_IS(4)) { // 0b0100 CALL_XATTN_KERNEL(half, false, true); @@ -57,4 +110,5 @@ extern "C" __global__ __aicore__ void x_attention(GM_ADDR query, GM_ADDR shared_ } else if (TILING_KEY_IS(10)) { // 0b1010 CALL_XATTN_KERNEL(bfloat16_t, true, false); } +#endif } From eb618ccf48abc0a66f3e01782c2a6eb0f9775770 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Tue, 18 Aug 2026 16:59:39 +0800 Subject: [PATCH 22/24] feat: slove the sparse_attn_sharekv_tiling bugs. --- .../op_host/sparse_attn_sharedkv_tiling.cpp | 7 ------- 1 file changed, 7 deletions(-) diff --git a/xllm_ops/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_tiling.cpp b/xllm_ops/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_tiling.cpp index decb775..5ab28e8 100644 --- a/xllm_ops/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_tiling.cpp +++ b/xllm_ops/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_tiling.cpp @@ -197,13 +197,6 @@ ge::graphStatus SASInfoParser::GetNpuInfo() 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; } From c64acc0747893998e500c104a5dea99d1c528da2 Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Tue, 25 Aug 2026 21:18:14 +0800 Subject: [PATCH 23/24] feat: change the catlass to cann --- .gitmodules | 2 +- third_party/catlass | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.gitmodules b/.gitmodules index 1762301..30fd4cb 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,6 +1,6 @@ [submodule "third_party/catlass"] path = third_party/catlass - url = https://gitcode.com/xLLM-AI/catlass.git + url = https://gitcode.com/cann/catlass.git [submodule "third_party/pto-isa"] path = third_party/pto-isa url = https://gitcode.com/cann/pto-isa.git diff --git a/third_party/catlass b/third_party/catlass index c02a6e8..dacc77c 160000 --- a/third_party/catlass +++ b/third_party/catlass @@ -1 +1 @@ -Subproject commit c02a6e8d9055d79601bee66e9dbfab24fffc41ee +Subproject commit dacc77c95d60355048d63a26d08e10b4e63a7367 From d9dc659e862272951c39f03e0aa7f70d60dabdec Mon Sep 17 00:00:00 2001 From: "ext.wangguangcai1" Date: Wed, 26 Aug 2026 15:02:08 +0800 Subject: [PATCH 24/24] feat: slove the cce compile error --- build.sh | 43 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/build.sh b/build.sh index ec43919..8efab5e 100755 --- a/build.sh +++ b/build.sh @@ -132,6 +132,46 @@ resolve_soc_version_list() { fi } +# ============================================================================ +# Catlass 兼容性补丁 +# +# 背景：catlass 的 block_epilogue_dequant.hpp 使用 `AscendC::DT_FLOAT` 等写法， +# 但 DT_FLOAT / DT_FLOAT16 / DT_BF16 在 CANN 的 kernel_type.h 中是预处理宏 +# (#define DT_FLOAT 0 ...)，宏无法被命名空间限定，`AscendC::DT_FLOAT` 会被 +# 预处理成 `AscendC::0`，编译报 "expected unqualified-id"。 +# 此处在构建前就地剥离 AscendC:: 前缀（仅限 build.log 确认的单文件），幂等。 +# ============================================================================ +patch_catlass_compat() { + local target="${BASE_DIR}/third_party/catlass/include/catlass/epilogue/block/block_epilogue_dequant.hpp" + if [[ ! -f "${target}" ]]; then + echo "[INFO] catlass patch: target not found, skip: ${target}" + return 0 + fi + + # 仅替换 build.log 确认报错的三种模式，避免扩大修改范围 + # 注意：grep -c 无匹配时退出码为 1 且仍输出 "0"，故用 `|| true` 兜底（勿用 echo 0， + # 否则会与 grep 的 "0" 拼成 "0\n0"，导致后续 [[ -eq ]] 解析失败）。 + local before + before=$(grep -cE 'AscendC::DT_(FLOAT16|BF16|FLOAT)' "${target}" 2>/dev/null || true) + before=${before:-0} + if [[ "${before}" -eq 0 ]]; then + echo "[INFO] catlass patch: already patched, skip" + return 0 + fi + + # 顺序：长串优先，避免 AscendC::DT_FLOAT 误吃 AscendC::DT_FLOAT16 + sed -i \ + -e 's/AscendC::DT_FLOAT16\b/DT_FLOAT16/g' \ + -e 's/AscendC::DT_BF16\b/DT_BF16/g' \ + -e 's/AscendC::DT_FLOAT\b/DT_FLOAT/g' \ + "${target}" + + local after + after=$(grep -cE 'AscendC::DT_(FLOAT16|BF16|FLOAT)' "${target}" 2>/dev/null || true) + after=${after:-0} + echo "[INFO] catlass patch: fixed $((before - after)) occurrence(s) in block_epilogue_dequant.hpp" +} + # ============================================================================ # 环境准备：设置编译器、清理打包产物 # ============================================================================ @@ -149,6 +189,9 @@ prepare_build_env() { $CC --version $CXX --version + # 修正 catlass 第三方依赖的宏不兼容问题（仅限单文件、幂等） + patch_catlass_compat + # 保留 BUILD_DIR 以支持增量编译，仅清理打包产物 rm -rf dist }