Xinjie-Q's picture
Upload Mage-VL: unified codec-native streaming VLM (image+video understanding + proactive gate)
12acbba verified
Raw
History Blame Contribute Delete
1.2 kB
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import os
import glob
import sys
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
cxx_flags = ["-O3"]
nvcc_flags = ["-O3", "--use_fast_math", "--extra-device-vectorization"]
# `-arch=native` detects the building machine's GPU at compile time, which is
# convenient interactively but (a) fails when no GPU is visible during the build
# and (b) can conflict with the `-gencode` flags torch appends from
# TORCH_CUDA_ARCH_LIST. So only use `native` when the caller has NOT pinned the
# arch via TORCH_CUDA_ARCH_LIST (e.g. cluster builds set it to "8.0").
if not os.environ.get("TORCH_CUDA_ARCH_LIST"):
nvcc_flags.append("-arch=native")
if sys.platform == 'win32':
cxx_flags = ["/O2"]
setup(
name='inference_extensions_cuda',
ext_modules=[
CUDAExtension(
name='inference_extensions_cuda',
sources=glob.glob('*.cpp') + glob.glob('*.cu'),
extra_compile_args={
"cxx": cxx_flags,
"nvcc": nvcc_flags,
},
),
],
cmdclass={
'build_ext': BuildExtension
}
)