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/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 } diff --git a/test/cpp_test/CMakeLists.txt b/test/cpp_test/CMakeLists.txt index 0e8ecc7..567cc31 100644 --- a/test/cpp_test/CMakeLists.txt +++ b/test/cpp_test/CMakeLists.txt @@ -83,21 +83,38 @@ 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") +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 - "${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}/${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" + "${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 +130,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 +172,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 +214,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 +221,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/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/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 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; } 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..a0136d7 100644 --- a/xllm_ops/build_aclnn.sh +++ b/xllm_ops/build_aclnn.sh @@ -323,13 +323,13 @@ 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" # ### JD's in-house operators #### "beam_search_group" - "x_attention" + "x_attention" "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/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 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); } 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; 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..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 @@ -186,12 +186,24 @@ __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. + static constexpr AscendC::ExpConfig MLA_EXP_CFG{AscendC::ExpAlgo::PRECISION_1ULP_FTZ_FALSE}; + 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 5bcb2ca..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 @@ -903,6 +903,20 @@ class MLAttentionDecoderAic { } 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], @@ -915,6 +929,7 @@ class MLAttentionDecoderAic { 0 // dstStride ); } +#endif SET_FLAG(MTE1, M, embed_split_idx % 2); @@ -976,14 +991,14 @@ class MLAttentionDecoderAic { 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 - ); + 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); } } @@ -991,6 +1006,9 @@ class MLAttentionDecoderAic { 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], @@ -1008,6 +1026,18 @@ class MLAttentionDecoderAic { 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], @@ -1020,6 +1050,7 @@ class MLAttentionDecoderAic { 0 // dstStride ); } +#endif SET_FLAG(MTE1, M, embed_split_idx % 2); @@ -1161,6 +1192,18 @@ class MLAttentionDecoderAic { 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], @@ -1173,6 +1216,7 @@ class MLAttentionDecoderAic { 0 // dstStride ); } +#endif } SET_FLAG(MTE1, MTE2, EVENT_ID7); } @@ -1342,6 +1386,18 @@ class MLAttentionDecoderAic { } 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], @@ -1354,6 +1410,7 @@ class MLAttentionDecoderAic { 0 // dstStride ); } +#endif SET_FLAG(MTE1, M, embed_split_idx % 2); if (embed_split_idx == 0 || embed_split_idx == 2) { @@ -1557,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], @@ -1569,6 +1638,7 @@ class MLAttentionDecoderAic { 0 // dstStride ); } +#endif SET_FLAG(MTE1, MTE2, l0_p_pingpong_flag + 6); SET_FLAG(MTE1, M, l0b_pingpong_flag); @@ -1792,6 +1862,10 @@ class MLADecoderAiv{ __aicore__ __attribute__((always_inline)) inline void Run() { + // 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); 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 } 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..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,11 +16,24 @@ 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" #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_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" 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 }