wehub-resource-sync
246 lines
8.7 KiB
Python
246 lines
8.7 KiB
Python
# Copyright (c) 2023 PaddlePaddle 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
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
from packaging.version import parse, Version
|
|
|
|
import paddle
|
|
from paddle.utils.cpp_extension import CUDAExtension, setup
|
|
|
|
sm_version = int(os.getenv("CUDA_SM_VERSION", "0"))
|
|
|
|
def get_nvcc_cuda_version(cuda_dir: str) -> Version:
|
|
"""Get the CUDA version from nvcc.
|
|
|
|
Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py
|
|
"""
|
|
nvcc_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"],
|
|
universal_newlines=True)
|
|
output = nvcc_output.split()
|
|
release_idx = output.index("release") + 1
|
|
nvcc_cuda_version = parse(output[release_idx].split(",")[0])
|
|
return nvcc_cuda_version
|
|
|
|
def update_git_submodule():
|
|
try:
|
|
subprocess.run(["git", "submodule", "update", "--init"], check=True)
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"Error occurred while updating git submodule: {str(e)}")
|
|
raise
|
|
|
|
|
|
def find_end_files(directory, end_str):
|
|
gen_files = []
|
|
for root, dirs, files in os.walk(directory):
|
|
for file in files:
|
|
if file.endswith(end_str):
|
|
gen_files.append(os.path.join(root, file))
|
|
return gen_files
|
|
|
|
|
|
def get_sm_version():
|
|
if sm_version > 0:
|
|
return sm_version
|
|
else:
|
|
prop = paddle.device.cuda.get_device_properties()
|
|
cc = prop.major * 10 + prop.minor
|
|
return cc
|
|
|
|
|
|
def strtobool(v):
|
|
if isinstance(v, bool):
|
|
return v
|
|
if v.lower() in ("yes", "true", "t", "y", "1"):
|
|
return True
|
|
elif v.lower() in ("no", "false", "f", "n", "0"):
|
|
return False
|
|
else:
|
|
raise ValueError(
|
|
f"Truthy value expected: got {v} but expected one of yes/no, true/false, t/f, y/n, 1/0 (case insensitive)."
|
|
)
|
|
|
|
|
|
def get_gencode_flags():
|
|
if not strtobool(os.getenv("FLAG_LLM_PDC", "False")):
|
|
cc = get_sm_version()
|
|
if cc == 90:
|
|
cc = f"{cc}a"
|
|
return ["-gencode", "arch=compute_{0},code=sm_{0}".format(cc)]
|
|
else:
|
|
# support more cuda archs
|
|
return [
|
|
"-gencode",
|
|
"arch=compute_80,code=sm_80",
|
|
"-gencode",
|
|
"arch=compute_75,code=sm_75",
|
|
"-gencode",
|
|
"arch=compute_70,code=sm_70",
|
|
]
|
|
|
|
|
|
gencode_flags = get_gencode_flags()
|
|
library_path = [os.environ.get("LD_LIBRARY_PATH", "/usr/local/cuda/lib64")]
|
|
|
|
sources = [
|
|
"./gpu/save_with_output.cc",
|
|
"./gpu/set_value_by_flags.cu",
|
|
"./gpu/token_penalty_multi_scores.cu",
|
|
"./gpu/token_penalty_multi_scores_v2.cu",
|
|
"./gpu/stop_generation_multi_ends.cu",
|
|
"./gpu/fused_get_rope.cu",
|
|
"./gpu/get_padding_offset.cu",
|
|
"./gpu/qkv_transpose_split.cu",
|
|
"./gpu/rebuild_padding.cu",
|
|
"./gpu/transpose_removing_padding.cu",
|
|
"./gpu/write_cache_kv.cu",
|
|
"./gpu/encode_rotary_qk.cu",
|
|
"./gpu/get_padding_offset_v2.cu",
|
|
"./gpu/rebuild_padding_v2.cu",
|
|
"./gpu/set_value_by_flags_v2.cu",
|
|
"./gpu/stop_generation_multi_ends_v2.cu",
|
|
"./gpu/get_output.cc",
|
|
"./gpu/save_with_output_msg.cc",
|
|
"./gpu/write_int8_cache_kv.cu",
|
|
"./gpu/step.cu",
|
|
"./gpu/quant_int8.cu",
|
|
"./gpu/dequant_int8.cu",
|
|
"./gpu/moe/preprocess_for_moe.cu",
|
|
"./gpu/get_position_ids_and_mask_encoder_batch.cu",
|
|
"./gpu/fused_rotary_position_encoding.cu",
|
|
"./gpu/flash_attn_bwd.cc",
|
|
"./gpu/tune_cublaslt_gemm.cu",
|
|
"./gpu/sample_kernels/top_p_sampling_reject.cu",
|
|
"./gpu/update_inputs_v2.cu",
|
|
"./gpu/noaux_tc.cu",
|
|
"./gpu/set_preids_token_penalty_multi_scores.cu",
|
|
"./gpu/speculate_decoding_kernels/ngram_match.cc",
|
|
"./gpu/speculate_decoding_kernels/speculate_save_output.cc",
|
|
"./gpu/speculate_decoding_kernels/speculate_get_output.cc",
|
|
"./gpu/save_output_dygraph.cu",
|
|
"./gpu/all_reduce.cu",
|
|
"./gpu/quantization/per_token_group_quant.cu",
|
|
"./gpu/quantization/per_tensor_quant_fp8.cu",
|
|
]
|
|
sources += find_end_files("./gpu/speculate_decoding_kernels", ".cu")
|
|
|
|
nvcc_compile_args = gencode_flags
|
|
update_git_submodule()
|
|
nvcc_compile_args += [
|
|
"-O3",
|
|
"-DNDEBUG",
|
|
"-U__CUDA_NO_HALF_OPERATORS__",
|
|
"-U__CUDA_NO_HALF_CONVERSIONS__",
|
|
"-U__CUDA_NO_BFLOAT16_OPERATORS__",
|
|
"-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
|
|
"-U__CUDA_NO_BFLOAT162_OPERATORS__",
|
|
"-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
|
|
]
|
|
|
|
include_dirs = [
|
|
"./gpu",
|
|
"./gpu/cutlass_kernels",
|
|
"./gpu/fp8_gemm_with_cutlass",
|
|
"./gpu/cutlass_kernels/fp8_gemm_fused/autogen",
|
|
"./third_party/cutlass/include",
|
|
"./third_party/cutlass/tools/util/include",
|
|
"./third_party/nlohmann_json/single_include",
|
|
"./gpu/sample_kernels",
|
|
"./gpu/moe/fused_moe",
|
|
]
|
|
cc = get_sm_version()
|
|
cuda_version = float(paddle.version.cuda())
|
|
nvcc_version = get_nvcc_cuda_version(os.environ.get("CUDA_HOME", "/usr/local/cuda"))
|
|
|
|
if cc >= 80:
|
|
sources += ["gpu/int8_gemm_with_cutlass/gemm_dequant.cu"]
|
|
|
|
sources += ["./gpu/append_attention.cu", "./gpu/multi_head_latent_attention.cu"]
|
|
|
|
sources += find_end_files("./gpu/append_attn", ".cu")
|
|
sources += find_end_files("./gpu/append_attn/template_instantiation", ".cu")
|
|
sources += find_end_files("./gpu/moe/fused_moe/cutlass_kernels/moe_gemm/", ".cu")
|
|
sources += find_end_files("./gpu/moe/fused_moe/", ".cu")
|
|
sources += "./gpu/cpp_extensions.cu",
|
|
|
|
|
|
fp8_auto_gen_directory = "gpu/cutlass_kernels/fp8_gemm_fused/autogen"
|
|
if os.path.isdir(fp8_auto_gen_directory):
|
|
shutil.rmtree(fp8_auto_gen_directory)
|
|
|
|
|
|
if cc == 89 and cuda_version >= 12.4:
|
|
os.system("python utils/auto_gen_fp8_fp8_gemm_fused_kernels.py --cuda_arch 89")
|
|
os.system("python utils/auto_gen_fp8_fp8_dual_gemm_fused_kernels.py --cuda_arch 89")
|
|
sources += find_end_files(fp8_auto_gen_directory, ".cu")
|
|
sources += [
|
|
"gpu/fp8_gemm_with_cutlass/fp8_fp8_half_gemm.cu",
|
|
"gpu/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.cu",
|
|
"gpu/fp8_gemm_with_cutlass/fp8_fp8_fp8_dual_gemm.cu",
|
|
]
|
|
|
|
if cc >= 80 and nvcc_version >= Version("12.4"):
|
|
os.environ.pop('PADDLE_CUDA_ARCH_LIST', None)
|
|
nvcc_compile_args += [
|
|
"-std=c++17",
|
|
"--use_fast_math",
|
|
"--threads=8",
|
|
"-D_GLIBCXX_USE_CXX11_ABI=1",
|
|
]
|
|
sources += ["./gpu/sage_attn_kernels/sageattn_fused.cu"]
|
|
if cc >= 80 and cc < 89:
|
|
sources += ["./gpu/sage_attn_kernels/sageattn_qk_int_sv_f16_kernel_sm80.cu"]
|
|
nvcc_compile_args += ["-gencode", "arch=compute_80,code=compute_80"]
|
|
elif cc >= 89 and cc < 90:
|
|
sources += ["./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_kernel_sm89.cu"]
|
|
nvcc_compile_args += ["-gencode", "arch=compute_89,code=compute_89"]
|
|
elif cc >= 90:
|
|
sources += [
|
|
"./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_kernel_sm90.cu",
|
|
"./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_dsk_kernel_sm90.cu",
|
|
]
|
|
nvcc_compile_args += ["-gencode", "arch=compute_90a,code=compute_90a"]
|
|
|
|
if cc >= 90 and cuda_version >= 12.0:
|
|
os.system("python utils/auto_gen_fp8_fp8_gemm_fused_kernels_sm90.py --cuda_arch 90")
|
|
os.system("python utils/auto_gen_fp8_fp8_gemm_fused_kernels_ptr_scale_sm90.py --cuda_arch 90")
|
|
os.system("python utils/auto_gen_fp8_fp8_dual_gemm_fused_kernels_sm90.py --cuda_arch 90")
|
|
os.system("python utils/auto_gen_fp8_fp8_block_gemm_fused_kernels_sm90.py --cuda_arch 90")
|
|
sources += find_end_files(fp8_auto_gen_directory, ".cu")
|
|
sources += [
|
|
"gpu/fp8_gemm_with_cutlass/fp8_fp8_half_gemm.cu",
|
|
"gpu/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.cu",
|
|
"gpu/fp8_gemm_with_cutlass/fp8_fp8_fp8_dual_gemm.cu",
|
|
"gpu/fp8_gemm_with_cutlass/fp8_fp8_half_block_gemm.cu",
|
|
"gpu/fp8_gemm_with_cutlass/fp8_fp8_half_gemm_ptr_scale.cu",
|
|
]
|
|
sources += find_end_files("./gpu/mla_attn", ".cu")
|
|
|
|
ops_name = f"paddlenlp_ops_{sm_version}" if sm_version != 0 else "paddlenlp_ops"
|
|
|
|
setup(
|
|
name=ops_name,
|
|
ext_modules=CUDAExtension(
|
|
sources=sources,
|
|
extra_compile_args={
|
|
"cxx": ["-O3", "-fopenmp", "-lgomp", "-std=c++17", "-DENABLE_BF16"],
|
|
"nvcc": nvcc_compile_args,
|
|
},
|
|
libraries=["cublasLt"],
|
|
library_dirs=library_path,
|
|
include_dirs=include_dirs,
|
|
),
|
|
)
|