Skip to content

Commit c80555a

Browse files
Infinite666Infinite666
committed
update gsa build scripts
1 parent 0e5d485 commit c80555a

6 files changed

Lines changed: 313 additions & 337 deletions

File tree

‎setup.py‎

Lines changed: 21 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ def print_platform_warning():
5151
{RED}{'=' * 80}
5252
{BOLD}⚠️ WARNING: PLATFORM environment variable is not set! ⚠️{RESET}
5353
{RED}{'=' * 80}{RESET}
54-
{YELLOW}Please set PLATFORM to one of: cuda, ascend, musa, maca{RESET}
54+
{YELLOW}Please set PLATFORM to one of: cuda, ascend, ascend-a3, musa, maca{RESET}
5555
Example:
5656
{BOLD}export PLATFORM=cuda{RESET} # For CUDA platform
5757
{YELLOW}In CI scenarios only, you don't need to specify PLATFORM. If it's not a CI scenario, please uninstall and then reinstall with PLATFORM specified.{RESET}
@@ -66,10 +66,18 @@ def print_platform_warning():
6666
atexit.register(print_platform_warning)
6767

6868

69+
def is_ascend() -> bool:
70+
return PLATFORM is not None and PLATFORM.startswith("ascend")
71+
72+
6973
def enable_sparse() -> bool:
7074
return ENABLE_SPARSE is not None and ENABLE_SPARSE.lower() == "true"
7175

7276

77+
def is_only_build_mode() -> bool:
78+
return "bdist_wheel" in sys.argv
79+
80+
7381
def is_editable_mode() -> bool:
7482
commands = [arg.lower() for arg in sys.argv]
7583
return (
@@ -94,21 +102,21 @@ def run(self):
94102
for ext in self.extensions:
95103
self.build_cmake(ext)
96104

97-
if enable_sparse() and PLATFORM == "ascend":
105+
if enable_sparse() and is_ascend():
106+
gsa_build_script = "ucm/sparse/gsa_on_device/csrc/ascend/build.sh"
107+
args = []
108+
if PLATFORM == "ascend-a3":
109+
args.append("a3")
110+
if not is_only_build_mode():
111+
args.append("install")
98112
try:
99-
print(
100-
"Running bash ucm/sparse/gsa_on_device/csrc/ascend/build_aclnn.sh to compiling NPU custom ops for UCM..."
101-
)
113+
print(f"Running {gsa_build_script} to compiling NPU custom ops for UCM...")
102114
subprocess.check_call(
103-
["bash", "ucm/sparse/gsa_on_device/csrc/ascend/build_aclnn.sh"]
104-
)
105-
print(
106-
"ucm/sparse/gsa_on_device/csrc/ascend/buid_aclnn.sh executed successfully!"
115+
["bash", gsa_build_script] + args
107116
)
117+
print(f"{gsa_build_script} executed successfully!")
108118
except subprocess.CalledProcessError as e:
109-
print(
110-
f"Error running ucm/sparse/gsa_on_device/csrc/ascend/build_aclnn.sh: {e}"
111-
)
119+
print("Error running {gsa_build_script}: {e}")
112120
raise SystemExit(e.returncode)
113121

114122
def build_cmake(self, ext: CMakeExtension):
@@ -129,7 +137,7 @@ def build_cmake(self, ext: CMakeExtension):
129137
match PLATFORM:
130138
case "cuda":
131139
cmake_args += ["-DRUNTIME_ENVIRONMENT=cuda"]
132-
case "ascend":
140+
case "ascend" | "ascend-a3":
133141
cmake_args += ["-DRUNTIME_ENVIRONMENT=ascend"]
134142
case "musa":
135143
cmake_args += ["-DRUNTIME_ENVIRONMENT=musa"]

‎ucm/sparse/gsa_on_device/csrc/ascend/build.sh‎

Lines changed: 76 additions & 175 deletions
Original file line numberDiff line numberDiff line change
@@ -1,194 +1,95 @@
11
#!/bin/bash
2-
# Copyright (c) 2024 Huawei Technologies Co., Ltd.
3-
# This file is a part of the CANN Open Software.
4-
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
5-
# Please refer to the License for details. You may not use this file except in compliance with the License.
6-
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7-
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8-
# See LICENSE in the root of the software repository for the full text of the License.
9-
# ======================================================================================================================
102

113
# SPDX-License-Identifier: Apache-2.0
124
# SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project
135
# Modified from
14-
# https://github.com/vllm-project/vllm-ascend/blob/main/csrc/build.sh
6+
# https://github.com/vllm-project/vllm-ascend/blob/main/csrc/build_aclnn.sh
157

16-
set -e
8+
ROOT_DIR=$(dirname $(readlink -f ${BASH_SOURCE[0]}))
179

18-
CURRENT_DIR=$(dirname $(readlink -f ${BASH_SOURCE[0]}))
19-
BUILD_DIR=${CURRENT_DIR}/build
20-
OUTPUT_DIR=${CURRENT_DIR}/output
21-
USER_ID=$(id -u)
22-
PARENT_JOB="false"
23-
CHECK_COMPATIBLE="true"
24-
ASAN="false"
25-
COV="false"
26-
VERBOSE="false"
10+
echo "ROOT_DIR: $ROOT_DIR"
2711

28-
if [ "${USER_ID}" != "0" ]; then
29-
DEFAULT_TOOLKIT_INSTALL_DIR="${HOME}/Ascend/ascend-toolkit/latest"
30-
DEFAULT_INSTALL_DIR="${HOME}/Ascend/latest"
31-
else
32-
DEFAULT_TOOLKIT_INSTALL_DIR="/usr/local/Ascend/ascend-toolkit/latest"
33-
DEFAULT_INSTALL_DIR="/usr/local/Ascend/latest"
34-
fi
35-
36-
CUSTOM_OPTION="-DBUILD_OPEN_PROJECT=ON"
37-
38-
function help_info() {
39-
echo "Usage: $0 [options]"
40-
echo "Options:"
41-
echo
42-
echo "-h|--help Displays help message."
43-
echo
44-
echo "-n|--op-name Specifies the compiled operator. If there are multiple values, separate them with semicolons and use quotation marks. The default is all."
45-
echo " For example: -n \"flash_attention_score\" or -n \"flash_attention_score;flash_attention_score_grad\""
46-
echo
47-
echo "-c|--compute-unit Specifies the chip type. If there are multiple values, separate them with semicolons and use quotation marks. The default is ascend910b."
48-
echo " For example: -c \"ascend910b\" or -c \"ascend910b;ascend310p\""
49-
echo
50-
echo "--cov Compiles with cov."
51-
echo
52-
echo "--verbose Displays more compilation information."
53-
echo
54-
}
55-
56-
function log() {
57-
local current_time=`date +"%Y-%m-%d %H:%M:%S"`
58-
echo "[$current_time] "$1
59-
}
60-
61-
function set_env()
62-
{
63-
source $ASCEND_CANN_PACKAGE_PATH/bin/setenv.bash || echo "0"
64-
65-
export BISHENG_REAL_PATH=$(which bisheng || true)
66-
67-
if [ -z "${BISHENG_REAL_PATH}" ];then
68-
log "Error: bisheng compilation tool not found, Please check whether the cann package or environment variables are set."
69-
exit 1
70-
fi
71-
}
72-
73-
function clean()
74-
{
75-
if [ -n "${BUILD_DIR}" ];then
76-
rm -rf ${BUILD_DIR}
77-
fi
78-
mkdir -p ${BUILD_DIR} ${OUTPUT_DIR}
79-
}
80-
81-
function cmake_config()
82-
{
83-
local extra_option="$1"
84-
log "Info: cmake config ${CUSTOM_OPTION} ${extra_option} ."
85-
cmake .. ${CUSTOM_OPTION} ${extra_option}
86-
}
87-
88-
function build()
89-
{
90-
local target="$1"
91-
if [ "${VERBOSE}" == "true" ];then
92-
local option="--verbose"
93-
fi
94-
cmake --build . --target ${target} ${JOB_NUM} ${option}
95-
}
96-
97-
function gen_bisheng(){
98-
local ccache_program=$1
99-
local gen_bisheng_dir=${BUILD_DIR}/gen_bisheng_dir
12+
IS_INSTALL="false"
13+
SOC_VERSION="ascend910b"
10014

101-
if [ ! -d "${gen_bisheng_dir}" ];then
102-
mkdir -p ${gen_bisheng_dir}
15+
for arg in "$@"; do
16+
if [[ "$arg" == "a3"]]; then
17+
SOC_VERSION="ascend910_93"
18+
elif [[ "$arg" == "install" ]]; then
19+
IS_INSTALL="true"
10320
fi
10421

105-
pushd ${gen_bisheng_dir}
106-
$(> bisheng)
107-
echo "#!/bin/bash" >> bisheng
108-
echo "ccache_args=""\"""${ccache_program} ${BISHENG_REAL_PATH}""\"" >> bisheng
109-
echo "args=""$""@" >> bisheng
110-
111-
if [ "${VERBOSE}" == "true" ];then
112-
echo "echo ""\"""$""{ccache_args} ""$""args""\"" >> bisheng
113-
fi
114-
115-
echo "eval ""\"""$""{ccache_args} ""$""args""\"" >> bisheng
116-
chmod +x bisheng
117-
118-
export PATH=${gen_bisheng_dir}:$PATH
119-
popd
120-
}
121-
122-
function build_package(){
123-
build package
124-
}
125-
126-
function build_host(){
127-
build_package
128-
}
129-
130-
function build_kernel(){
131-
build ops_kernel
132-
}
133-
134-
while [[ $# -gt 0 ]]; do
135-
case $1 in
136-
-h|--help)
137-
help_info
138-
exit
139-
;;
140-
-n|--op-name)
141-
ascend_op_name="$2"
142-
shift 2
143-
;;
144-
-c|--compute-unit)
145-
ascend_compute_unit="$2"
146-
shift 2
147-
;;
148-
*)
149-
help_info
22+
echo "SOC_VERSION: $SOC_VERSION"
23+
24+
25+
if [[ "$SOC_VERSION" =~ ^ascend310 ]]; then
26+
# ASCEND310P series
27+
# currently, no custom aclnn ops for ASCEND310 series
28+
# CUSTOM_OPS=""
29+
# SOC_ARG="ascend310p"
30+
exit 0
31+
elif [[ "$SOC_VERSION" =~ ^ascend910b ]]; then
32+
# ASCEND910B (A2) series
33+
# dependency: catlass
34+
CUSTOM_OPS="hamming_dist_top_k;reshape_and_cache_bnsd;"
35+
SOC_ARG="ascend910b"
36+
elif [[ "$SOC_VERSION" =~ ^ascend910_93 ]]; then
37+
# ASCEND910C (A3) series
38+
# dependency: catlass
39+
# dependency: cann-toolkit file moe_distribute_base.h
40+
HCCL_STRUCT_FILE_PATH=$(find -L "${ASCEND_TOOLKIT_HOME}" -name "moe_distribute_base.h" 2>/dev/null | head -n1)
41+
if [ -z "$HCCL_STRUCT_FILE_PATH" ]; then
42+
echo "cannot find moe_distribute_base.h file in CANN env"
15043
exit 1
151-
;;
152-
esac
153-
done
154-
155-
if [ -n "${ascend_compute_unit}" ];then
156-
CUSTOM_OPTION="${CUSTOM_OPTION} -DASCEND_COMPUTE_UNIT=${ascend_compute_unit}"
157-
fi
158-
159-
if [ -n "${ascend_op_name}" ];then
160-
CUSTOM_OPTION="${CUSTOM_OPTION} -DASCEND_OP_NAME=${ascend_op_name}"
161-
fi
162-
163-
if [ -n "${ASCEND_HOME_PATH}" ];then
164-
ASCEND_CANN_PACKAGE_PATH=${ASCEND_HOME_PATH}
165-
elif [ -n "${ASCEND_OPP_PATH}" ];then
166-
ASCEND_CANN_PACKAGE_PATH=$(dirname ${ASCEND_OPP_PATH})
167-
elif [ -d "${DEFAULT_TOOLKIT_INSTALL_DIR}" ];then
168-
ASCEND_CANN_PACKAGE_PATH=${DEFAULT_TOOLKIT_INSTALL_DIR}
169-
elif [ -d "${DEFAULT_INSTALL_DIR}" ];then
170-
ASCEND_CANN_PACKAGE_PATH=${DEFAULT_INSTALL_DIR}
44+
fi
45+
# for dispatch_gmm_combine_decode
46+
yes | cp "${HCCL_STRUCT_FILE_PATH}" "${ROOT_DIR}/csrc/utils/inc/kernel"
47+
# for dispatch_ffn_combine
48+
SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)
49+
TARGET_DIR="$SCRIPT_DIR/dispatch_ffn_combine/op_kernel/utils/"
50+
TARGET_FILE="$TARGET_DIR/$(basename "$HCCL_STRUCT_FILE_PATH")"
51+
# for dispatch_ffn_combine_bf16
52+
SCRIPT_DIR_BF16=$(cd "$(dirname "$0")" && pwd)
53+
TARGET_DIR_BF16="$SCRIPT_DIR_BF16/dispatch_ffn_combine_bf16/op_kernel/utils/"
54+
TARGET_FILE_BF16="$TARGET_DIR_BF16/$(basename "$HCCL_STRUCT_FILE_PATH")"
55+
56+
echo "*************************************"
57+
echo $HCCL_STRUCT_FILE_PATH
58+
echo "$TARGET_DIR"
59+
cp "$HCCL_STRUCT_FILE_PATH" "$TARGET_DIR"
60+
cp "$HCCL_STRUCT_FILE_PATH" "$TARGET_DIR_BF16"
61+
62+
sed -i 's/struct HcclOpResParam {/struct HcclOpResParamCustom {/g' "$TARGET_FILE"
63+
sed -i 's/struct HcclRankRelationResV2 {/struct HcclRankRelationResV2Custom {/g' "$TARGET_FILE"
64+
sed -i 's/struct HcclOpResParam {/struct HcclOpResParamCustom {/g' "$TARGET_FILE_BF16"
65+
sed -i 's/struct HcclRankRelationResV2 {/struct HcclRankRelationResV2Custom {/g' "$TARGET_FILE_BF16"
66+
67+
CUSTOM_OPS_ARRAY=(
68+
"hamming_dist_top_k"
69+
"reshape_and_cache_bnsd"
70+
)
71+
CUSTOM_OPS=$(IFS=';'; echo "${CUSTOM_OPS_ARRAY[*]}")
72+
SOC_ARG="ascend910_93"
17173
else
172-
log "Error: Please set the toolkit package installation directory through parameter -p|--package-path."
173-
exit 1
74+
# others
75+
# currently, no custom aclnn ops for other series
76+
exit 0
17477
fi
17578

176-
if [ "${PARENT_JOB}" == "false" ];then
177-
CPU_NUM=$(($(cat /proc/cpuinfo | grep "^processor" | wc -l)*2))
178-
JOB_NUM="-j${CPU_NUM}"
179-
fi
18079

181-
CUSTOM_OPTION="${CUSTOM_OPTION} -DCUSTOM_ASCEND_CANN_PACKAGE_PATH=${ASCEND_CANN_PACKAGE_PATH} -DCHECK_COMPATIBLE=${CHECK_COMPATIBLE}"
80+
# build custom ops
81+
cd $ROOT_DIR
82+
echo "Current directory: $PWD, begin to build custom ops (ascendc version)..."
83+
rm -rf build output
84+
echo "Building custom ops $CUSTOM_OPS for $SOC_VERSION"
85+
bash build_ops.sh -n "$CUSTOM_OPS" -c "$SOC_ARG"
18286

183-
set_env
184-
clean
87+
# build python wheel
88+
echo "Begin to build python wheel..."
89+
bash build_wheel.sh
18590

186-
ccache_system=$(which ccache || true)
187-
if [ -n "${ccache_system}" ];then
188-
CUSTOM_OPTION="${CUSTOM_OPTION} -DENABLE_CCACHE=ON -DCUSTOM_CCACHE=${ccache_system}"
189-
gen_bisheng ${ccache_system}
91+
# install custom ops and python package
92+
echo "Installing ucm custom ops. This may take a while, please wait..."
93+
if [[ "$IS_INSTALL" == "true" ]]; then
94+
bash install.sh
19095
fi
191-
192-
cd ${BUILD_DIR}
193-
cmake_config
194-
build_package

0 commit comments

Comments
 (0)