初始化项目,由ModelHub XC社区提供模型
Model: ayh015/myLightningOPD Source: Original Platform
This commit is contained in:
42
slime/backends/megatron_utils/kernels/int4_qat/setup.py
Normal file
42
slime/backends/megatron_utils/kernels/int4_qat/setup.py
Normal file
@@ -0,0 +1,42 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from setuptools import setup
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
import torch
|
||||
|
||||
# Get CUDA arch list
|
||||
arch_list = []
|
||||
if torch.cuda.is_available():
|
||||
for i in range(torch.cuda.device_count()):
|
||||
major, minor = torch.cuda.get_device_capability(i)
|
||||
arch_list.append(f"{major}.{minor}")
|
||||
arch_list = sorted(set(arch_list))
|
||||
|
||||
setup(
|
||||
name="fake_int4_quant_cuda",
|
||||
ext_modules=[
|
||||
CUDAExtension(
|
||||
name="fake_int4_quant_cuda",
|
||||
sources=["fake_int4_quant_cuda.cu"],
|
||||
extra_compile_args={
|
||||
"cxx": [
|
||||
"-O3",
|
||||
"-std=c++17",
|
||||
],
|
||||
"nvcc": [
|
||||
"-O3",
|
||||
"-std=c++17",
|
||||
"--expt-relaxed-constexpr",
|
||||
"-Xcompiler",
|
||||
"-fPIC",
|
||||
]
|
||||
+ [
|
||||
f'-gencode=arch=compute_{arch.replace(".", "")},code=sm_{arch.replace(".", "")}'
|
||||
for arch in arch_list
|
||||
],
|
||||
},
|
||||
)
|
||||
],
|
||||
cmdclass={"build_ext": BuildExtension},
|
||||
)
|
||||
Reference in New Issue
Block a user