feat: adapter the a5

.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

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
 }

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,15 +214,13 @@ 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)
 
 # 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)

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,

test/cpp_test/utils_tensor.h

@@ -93,11 +93,23 @@ int Init(int32_t deviceId, aclrtStream* stream) {
 std::vector<int64_t> get_weight_storage_shape(const std::vector<int64_t>& shape)
 {
     std::vector<int64_t> 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;
 }
 

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)
+    torch.testing.assert_close(out, golden, atol=6e-2, rtol=6e-2)

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;
 }
 

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}/../../../

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"