setup.py 24.3 KB
Newer Older
chenxl's avatar
chenxl committed
1
2
3
#!/usr/bin/env python
# coding=utf-8
'''
Xiaodong Ye's avatar
Xiaodong Ye committed
4
Description  :
chenxl's avatar
chenxl committed
5
Author       : chenxl
chenxl's avatar
chenxl committed
6
Date         : 2024-07-27 16:15:27
chenxl's avatar
chenxl committed
7
Version      : 1.0.0
Xiaodong Ye's avatar
Xiaodong Ye committed
8
LastEditors  : chenxl
9
LastEditTime : 2024-08-14 16:36:19
chenxl's avatar
chenxl committed
10
11
12
Adapted from:
https://github.com/Dao-AILab/flash-attention/blob/v2.6.3/setup.py
Copyright (c) 2023, Tri Dao.
Xiaodong Ye's avatar
Xiaodong Ye committed
13
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
chenxl's avatar
chenxl committed
14
'''
chenxl's avatar
chenxl committed
15

chenxl's avatar
chenxl committed
16
17
18
19
import os
import sys
import re
import ast
20
from collections import deque
chenxl's avatar
chenxl committed
21
import subprocess
22
23
import select
import time
chenxl's avatar
chenxl committed
24
import platform
chenxl's avatar
chenxl committed
25
import shutil
26
from typing import List, Optional, Literal
27
import http.client
chenxl's avatar
chenxl committed
28
29
import urllib.request
import urllib.error
chenxl's avatar
chenxl committed
30
31
from pathlib import Path
from packaging.version import parse
32
import torch
chenxl's avatar
chenxl committed
33
34
35
import torch.version
from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
from setuptools import setup, Extension
Azure-Tang's avatar
Azure-Tang committed
36
from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME, ROCM_HOME
Xiaodong Ye's avatar
Xiaodong Ye committed
37
38
39
40
41
try:
    from torch_musa.utils.simple_porting import SimplePorting
    from torch_musa.utils.musa_extension import BuildExtension, MUSAExtension, MUSA_HOME
except ImportError:
    MUSA_HOME=None
42
43
    
with_balance = os.environ.get("USE_BALANCE_SERVE", "0") == "1"
chenxl's avatar
chenxl committed
44

45
46
47
48
49
50
51
52
53
class CpuInstructInfo:
    CPU_INSTRUCT = os.getenv("CPU_INSTRUCT", "NATIVE")
    FANCY = "FANCY"
    AVX512 = "AVX512"
    AVX2 = "AVX2"
    CMAKE_NATIVE = "-DLLAMA_NATIVE=ON"
    CMAKE_FANCY = "-DLLAMA_NATIVE=OFF -DLLAMA_FMA=ON -DLLAMA_F16C=ON -DLLAMA_AVX=ON -DLLAMA_AVX2=ON -DLLAMA_AVX512=ON -DLLAMA_AVX512_FANCY_SIMD=ON"
    CMAKE_AVX512 = "-DLLAMA_NATIVE=OFF -DLLAMA_FMA=ON -DLLAMA_F16C=ON -DLLAMA_AVX=ON -DLLAMA_AVX2=ON -DLLAMA_AVX512=ON"
    CMAKE_AVX2 = "-DLLAMA_NATIVE=OFF -DLLAMA_FMA=ON -DLLAMA_F16C=ON -DLLAMA_AVX=ON -DLLAMA_AVX2=ON"
Xiaodong Ye's avatar
Xiaodong Ye committed
54

chenxl's avatar
chenxl committed
55
56
57
class VersionInfo:
    THIS_DIR = os.path.dirname(os.path.abspath(__file__))
    PACKAGE_NAME = "ktransformers"
chenxl's avatar
chenxl committed
58
59
60
61
62
    BASE_WHEEL_URL:str = (
        "https://github.com/kvcache-ai/ktransformers/releases/download/{tag_name}/{wheel_filename}"
    )
    FORCE_BUILD = os.getenv("KTRANSFORMERS_FORCE_BUILD", "FALSE") == "TRUE"

Xiaodong Ye's avatar
Xiaodong Ye committed
63
64
65
66
67
68
69
70
71
72
    def get_musa_bare_metal_version(self, musa_dir):
        raw_output = subprocess.run(
            [musa_dir + "/bin/mcc", "-v"], check=True,
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT).stdout.decode("utf-8")
        output = raw_output.split()
        release_idx = output.index("version") + 1
        bare_metal_version = parse(output[release_idx].split(",")[0])
        musa_version = f"{bare_metal_version.major}{bare_metal_version.minor}"
        return musa_version

Azure-Tang's avatar
Azure-Tang committed
73
74
75
    def get_rocm_bare_metal_version(self, rocm_dir):
        """
        Get the ROCm version from the ROCm installation directory.
76

Azure-Tang's avatar
Azure-Tang committed
77
78
        Args:
            rocm_dir: Path to the ROCm installation directory
79

Azure-Tang's avatar
Azure-Tang committed
80
81
82
83
84
85
        Returns:
            A string representation of the ROCm version (e.g., "63" for ROCm 6.3)
        """
        try:
            # Try using rocm_agent_enumerator to get version info
            raw_output = subprocess.check_output(
86
                [rocm_dir + "/bin/rocminfo", "--version"],
Azure-Tang's avatar
Azure-Tang committed
87
88
89
90
91
92
93
94
95
96
97
98
                universal_newlines=True,
                stderr=subprocess.STDOUT)
            # Extract version number from output
            match = re.search(r'(\d+\.\d+)', raw_output)
            if match:
                version_str = match.group(1)
                version = parse(version_str)
                rocm_version = f"{version.major}{version.minor}"
                return rocm_version
        except (subprocess.CalledProcessError, FileNotFoundError):
            # If rocminfo --version fails, try alternative methods
            pass
99

Azure-Tang's avatar
Azure-Tang committed
100
101
102
103
104
105
106
107
108
        try:
            # Try reading version from release file
            with open(os.path.join(rocm_dir, "share/doc/hip/version.txt"), "r") as f:
                version_str = f.read().strip()
                version = parse(version_str)
                rocm_version = f"{version.major}{version.minor}"
                return rocm_version
        except (FileNotFoundError, IOError):
            pass
109

Azure-Tang's avatar
Azure-Tang committed
110
111
112
113
114
115
116
117
        # If all else fails, try to extract from directory name
        dir_name = os.path.basename(os.path.normpath(rocm_dir))
        match = re.search(r'rocm-(\d+\.\d+)', dir_name)
        if match:
            version_str = match.group(1)
            version = parse(version_str)
            rocm_version = f"{version.major}{version.minor}"
            return rocm_version
118

Azure-Tang's avatar
Azure-Tang committed
119
120
121
122
123
124
125
126
127
128
129
130
131
132
        # Fallback to extracting from hipcc version
        try:
            raw_output = subprocess.check_output(
                [rocm_dir + "/bin/hipcc", "--version"],
                universal_newlines=True,
                stderr=subprocess.STDOUT)
            match = re.search(r'HIP version: (\d+\.\d+)', raw_output)
            if match:
                version_str = match.group(1)
                version = parse(version_str)
                rocm_version = f"{version.major}{version.minor}"
                return rocm_version
        except (subprocess.CalledProcessError, FileNotFoundError):
            pass
133

Azure-Tang's avatar
Azure-Tang committed
134
135
136
        # If we still can't determine the version, raise an error
        raise ValueError(f"Could not determine ROCm version from directory: {rocm_dir}")

chenxl's avatar
chenxl committed
137
    def get_cuda_bare_metal_version(self, cuda_dir):
chenxl's avatar
chenxl committed
138
139
        raw_output = subprocess.check_output(
            [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True)
chenxl's avatar
chenxl committed
140
141
142
143
144
        output = raw_output.split()
        release_idx = output.index("release") + 1
        bare_metal_version = parse(output[release_idx].split(",")[0])
        cuda_version = f"{bare_metal_version.major}{bare_metal_version.minor}"
        return cuda_version
chenxl's avatar
chenxl committed
145

Xiaodong Ye's avatar
Xiaodong Ye committed
146
    def get_cuda_version_of_torch(self):
chenxl's avatar
chenxl committed
147
148
149
        torch_cuda_version = parse(torch.version.cuda)
        cuda_version = f"{torch_cuda_version.major}{torch_cuda_version.minor}"
        return cuda_version
chenxl's avatar
chenxl committed
150

chenxl's avatar
chenxl committed
151
152
153
154
155
156
    def get_platform(self,):
        """
        Returns the platform name as used in wheel filenames.
        """
        if sys.platform.startswith("linux"):
            return f'linux_{platform.uname().machine}'
chenxl's avatar
chenxl committed
157
158
        elif sys.platform == "win32":
            return "win_amd64"
chenxl's avatar
chenxl committed
159
160
        else:
            raise ValueError("Unsupported platform: {}".format(sys.platform))
chenxl's avatar
chenxl committed
161

chenxl's avatar
chenxl committed
162
    def get_cpu_instruct(self,):
163
164
165
166
167
168
169
170
        if CpuInstructInfo.CPU_INSTRUCT == CpuInstructInfo.FANCY:
            return "fancy"
        elif CpuInstructInfo.CPU_INSTRUCT == CpuInstructInfo.AVX512:
            return "avx512"
        elif CpuInstructInfo.CPU_INSTRUCT == CpuInstructInfo.AVX2:
            return "avx2"
        else:
            print("Using native cpu instruct")
chenxl's avatar
chenxl committed
171
        if sys.platform.startswith("linux"):
chenxl's avatar
chenxl committed
172
            with open('/proc/cpuinfo', 'r', encoding="utf-8") as cpu_f:
chenxl's avatar
chenxl committed
173
                cpuinfo = cpu_f.read()
chenxl's avatar
chenxl committed
174
175
            flags_line = [line for line in cpuinfo.split(
                '\n') if line.startswith('flags')][0]
chenxl's avatar
chenxl committed
176
            flags = flags_line.split(':')[1].strip().split(' ')
177
178
179
180
            # fancy with AVX512-VL, AVX512-BW, AVX512-DQ, AVX512-VNNI
            for flag in flags:
                if 'avx512bw' in flag:
                    return 'fancy'
chenxl's avatar
chenxl committed
181
182
183
184
185
186
            for flag in flags:
                if 'avx512' in flag:
                    return 'avx512'
            for flag in flags:
                if 'avx2' in flag:
                    return 'avx2'
chenxl's avatar
chenxl committed
187
188
            raise ValueError(
                "Unsupported cpu Instructions: {}".format(flags_line))
chenxl's avatar
chenxl committed
189
        elif sys.platform == "win32":
jzl's avatar
jzl committed
190
191
            from cpufeature.extension import CPUFeature

chenxl's avatar
chenxl committed
192
193
194
195
196
197
198
199
            if CPUFeature.get("AVX512bw", False):
                return 'fancy'
            if CPUFeature.get("AVX512f", False):
                return 'avx512'
            if CPUFeature.get("AVX2", False):
                return 'avx2'
            raise ValueError(
                "Unsupported cpu Instructions: {}".format(str(CPUFeature)))
chenxl's avatar
chenxl committed
200
201
202
        else:
            raise ValueError("Unsupported platform: {}".format(sys.platform))

chenxl's avatar
chenxl committed
203
204
205
206
    def get_torch_version(self,):
        torch_version_raw = parse(torch.__version__)
        torch_version = f"{torch_version_raw.major}{torch_version_raw.minor}"
        return torch_version
Xiaodong Ye's avatar
Xiaodong Ye committed
207

chenxl's avatar
chenxl committed
208
209
210
    def get_flash_version(self,):
        version_file = os.path.join(
            Path(VersionInfo.THIS_DIR), VersionInfo.PACKAGE_NAME, "__init__.py")
chenxl's avatar
chenxl committed
211
        with open(version_file, "r", encoding="utf-8") as f:
chenxl's avatar
chenxl committed
212
213
214
215
216
217
            version_match = re.search(
                r"^__version__\s*=\s*(.*)$", f.read(), re.MULTILINE)
        flash_version = ast.literal_eval(version_match.group(1))
        return flash_version

    def get_package_version(self, full_version=False):
Xiaodong Ye's avatar
Xiaodong Ye committed
218
219
220
221
222
        flash_version = str(self.get_flash_version())
        torch_version = self.get_torch_version()
        cpu_instruct = self.get_cpu_instruct()
        backend_version = ""
        if CUDA_HOME is not None:
223
            backend_version = f"cu{self.get_cuda_bare_metal_version(CUDA_HOME)}"
Xiaodong Ye's avatar
Xiaodong Ye committed
224
225
        elif MUSA_HOME is not None:
            backend_version = f"mu{self.get_musa_bare_metal_version(MUSA_HOME)}"
Azure-Tang's avatar
Azure-Tang committed
226
227
        elif ROCM_HOME is not None:
            backend_version = f"rocm{self.get_rocm_bare_metal_version(ROCM_HOME)}"
Xiaodong Ye's avatar
Xiaodong Ye committed
228
        else:
Azure-Tang's avatar
Azure-Tang committed
229
            raise ValueError("Unsupported backend: CUDA_HOME MUSA_HOME ROCM_HOME all not set.")
Xiaodong Ye's avatar
Xiaodong Ye committed
230
        package_version = f"{flash_version}+{backend_version}torch{torch_version}{cpu_instruct}"
chenxl's avatar
chenxl committed
231
232
233
        if full_version:
            return package_version
        if not VersionInfo.FORCE_BUILD:
Xiaodong Ye's avatar
Xiaodong Ye committed
234
            return flash_version
chenxl's avatar
chenxl committed
235
        return package_version
chenxl's avatar
chenxl committed
236

chenxl's avatar
chenxl committed
237
238
239
240

class BuildWheelsCommand(_bdist_wheel):
    def get_wheel_name(self,):
        version_info = VersionInfo()
chenxl's avatar
chenxl committed
241
242
        package_version = version_info.get_package_version(full_version=True)
        flash_version = version_info.get_flash_version()
chenxl's avatar
chenxl committed
243
        python_version = f"cp{sys.version_info.major}{sys.version_info.minor}"
chenxl's avatar
chenxl committed
244
245
246
247
248
        wheel_filename = f"{VersionInfo.PACKAGE_NAME}-{package_version}-{python_version}-{python_version}-{version_info.get_platform()}.whl"
        wheel_url = VersionInfo.BASE_WHEEL_URL.format(tag_name=f"v{flash_version}", wheel_filename=wheel_filename)
        return wheel_filename, wheel_url


chenxl's avatar
chenxl committed
249
    def run(self):
chenxl's avatar
chenxl committed
250
251
        if VersionInfo.FORCE_BUILD:
            super().run()
252
            return
chenxl's avatar
chenxl committed
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
        wheel_filename, wheel_url = self.get_wheel_name()
        print("Guessing wheel URL: ", wheel_url)
        try:
            urllib.request.urlretrieve(wheel_url, wheel_filename)
            # Make the archive
            # Lifted from the root wheel processing command
            # https://github.com/pypa/wheel/blob/cf71108ff9f6ffc36978069acb28824b44ae028e/src/wheel/bdist_wheel.py#LL381C9-L381C85
            if not os.path.exists(self.dist_dir):
                os.makedirs(self.dist_dir)

            impl_tag, abi_tag, plat_tag = self.get_tag()
            archive_basename = f"{self.wheel_dist_name}-{impl_tag}-{abi_tag}-{plat_tag}"

            wheel_path = os.path.join(self.dist_dir, archive_basename + ".whl")
            print("Raw wheel path", wheel_path)
chenxl's avatar
chenxl committed
268
            shutil.move(wheel_filename, wheel_path)
269
        except (urllib.error.HTTPError, urllib.error.URLError, http.client.RemoteDisconnected):
chenxl's avatar
chenxl committed
270
271
272
273
            print("Precompiled wheel not found. Building from source...")
            # If the wheel could not be downloaded, build from source
            super().run()

chenxl's avatar
chenxl committed
274

275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
ANSI_ESCAPE = re.compile(
    r'\033[@-Z\\-_\[\]P]|\033\[[0-?]*[ -/]*[@-~]|\033][^\007\033]*\007|[\000-\037]'
)

def colored(text, color=None, bold=False):
    fmt = []
    if color== 'red':
        fmt.append('31')
    elif color == 'green':
        fmt.append('32')
    if bold:
        fmt.append('1')

    return f"\033[{';'.join(fmt)}m{text}\033[0m"


def split_line(text: str) -> List[str]:
    """Split text into lines based on terminal width."""
    term_width = shutil.get_terminal_size().columns or 80
    if not text.strip():
        return []
    # Split by explicit newlines and wrap long lines
    lines = []
    for line in text.split('\n'):
        while len(line) > term_width:
            lines.append(line[:term_width])
            line = line[term_width:]
        if line:
            lines.append(line)
    return lines



ANSI_ESCAPE = re.compile(
    r'\033[@-Z\\-_\[\]P]|\033\[[0-?]*[ -/]*[@-~]|\033][^\007\033]*\007|[\000-\037]'
)

def colored(text, color=None, bold=False):
    fmt = []
    if color== 'red':
        fmt.append('31')
    elif color == 'green':
        fmt.append('32')
    if bold:
        fmt.append('1')

    return f"\033[{';'.join(fmt)}m{text}\033[0m"


def split_line(text: str) -> List[str]:
    """Split text into lines based on terminal width."""
    term_width = shutil.get_terminal_size().columns or 80
    if not text.strip():
        return []
    # Split by explicit newlines and wrap long lines
    lines = []
    for line in text.split('\n'):
        while len(line) > term_width:
            lines.append(line[:term_width])
            line = line[term_width:]
        if line:
            lines.append(line)
    return lines


def run_command_with_live_tail(ext: str, command: List[str], output_lines: int = 20,
                               refresh_rate: float = 0.1, cwd: Optional[str] = None):
    """
    Execute a script-like command with real-time output of the last `output_lines` lines.

    - during execution: displays the last `output_lines` lines of output in real-time.
    - On success: Clears the displayed output.
    - On failure: Prints the full command output.

    Args:
        ext (str): the name of the native extension currently building.
        command (List[str]): The command to execute, as a list of arguments.
        output_lines (int, optional): Number of terminal lines to display during live output. Defaults to 20.
        refresh_rate (float, optional): Time in seconds between output refreshes. Defaults to 0.1.
        cwd (Optional[str], optional): Working directory to run the command in. Defaults to current directory.
    """
    # Dump all subprocess output without any buffering if stdout is not a terminal
    if not sys.stdout.isatty():
        return subprocess.run(command, cwd=cwd, check=True)
    # Start time for elapsed time calculation
    start = time.time()
    # Buffer for all output
    all_output = []
    write_buffer = deque(maxlen=output_lines)
    # Current number of lines from sub process displayed
    current_lines = 0

    # ANSI escape codes for terminal control
    CLEAR_LINE = '\033[K'
    MOVE_UP = '\033[1A'
    SAVE_CURSOR = '\0337'
    RESTORE_CURSOR = '\0338'
    CLEAR_REMAINING = '\033[J'

    def write_progress(status: Literal['RUNNING', 'SUCCEED', 'FAILED'] = 'RUNNING',
                       new_line: Optional[str] = None):
        """Update terminal display with latest output"""
        nonlocal current_lines, process
        sys.stdout.write(SAVE_CURSOR)
        sys.stdout.write(MOVE_UP * current_lines)
        banner = f"ext={ext} pid={process.pid} status={status.upper()} elapsed=({time.time()-start:.2f}S)\n"
        if status != 'FAILED':
            banner = colored(banner, 'green', bold=True)
        else:
            banner = colored(banner, 'red', bold=True)
        sys.stdout.write(CLEAR_LINE + banner)
        if new_line is not None:
            all_output.append(new_line)
            write_buffer.extend(split_line(ANSI_ESCAPE.sub('', new_line).rstrip()))
        elif status == 'RUNNING':
            sys.stdout.write(RESTORE_CURSOR)
            sys.stdout.flush()
            return

        sys.stdout.write(CLEAR_REMAINING)
        if status == 'RUNNING':
            current_lines = 1 + len(write_buffer)
            for text in write_buffer:
                sys.stdout.write(text + '\n')
        elif status == 'FAILED':
            for text in all_output:
                sys.stdout.write(text)
        sys.stdout.flush()

    # Start subprocess
    sys.stdout.write(colored(f'ext={ext} command={" ".join(str(c) for c in command)}\n', bold=True))
    sys.stdout.flush()
    process = subprocess.Popen(
        command,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        cwd=cwd,
        text=True,
        bufsize=1
    )

    try:
        write_progress()
        poll_obj = select.poll()
        poll_obj.register(process.stdout, select.POLLIN)
        while process.poll() is None:
            poll_result = poll_obj.poll(refresh_rate * 1000)
            if poll_result:
                write_progress(new_line=process.stdout.readline())
            else:
                write_progress()

        # Get any remaining output
        while True:
            line = process.stdout.readline()
            if not line:
                break
            write_progress(new_line=line)
    except BaseException as e:
        process.terminate()
        raise e
    finally:
        exit_code = process.wait()
        write_progress(status='SUCCEED' if exit_code == 0 else 'FAILED')


chenxl's avatar
chenxl committed
441
442
443
444
445
446
447
448
# Convert distutils Windows platform specifiers to CMake -A arguments
PLAT_TO_CMAKE = {
    "win32": "Win32",
    "win-amd64": "x64",
    "win-arm32": "ARM",
    "win-arm64": "ARM64",
}

chenxl's avatar
chenxl committed
449

chenxl's avatar
chenxl committed
450
class CMakeExtension(Extension):
451
    def __init__(self, name: str, sourcedir: str) -> None:
chenxl's avatar
chenxl committed
452
        super().__init__(name, sources=[])
453
454
        print(name, sourcedir)
        self.sourcedir = sourcedir
chenxl's avatar
chenxl committed
455

456
457
458
459
460
461
462
def get_cmake_abi_args(cmake_args):
    if torch.compiled_with_cxx11_abi():
        cmake_args.append("-D_GLIBCXX_USE_CXX11_ABI=1")
    else:
        cmake_args.append("-D_GLIBCXX_USE_CXX11_ABI=0")
    return cmake_args

chenxl's avatar
chenxl committed
463
class CMakeBuild(BuildExtension):
chenxl's avatar
chenxl committed
464

chenxl's avatar
chenxl committed
465
466
467
468
469
470
471
472
473
474
    def build_extension(self, ext) -> None:
        if not isinstance(ext, CMakeExtension):
            super().build_extension(ext)
            return
        ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name)
        extdir = ext_fullpath.parent.resolve()

        # Using this requires trailing slash for auto-detection & inclusion of
        # auxiliary "native" libs

chenxl's avatar
chenxl committed
475
476
        debug = int(os.environ.get("DEBUG", 0)
                    ) if self.debug is None else self.debug
chenxl's avatar
chenxl committed
477
478
479
480
481
482
483
484
485
486
487
488
489
490
        cfg = "Debug" if debug else "Release"

        # CMake lets you override the generator - we need to check this.
        # Can be set with Conda-Build, for example.
        cmake_generator = os.environ.get("CMAKE_GENERATOR", "")

        # Set Python_EXECUTABLE instead if you use PYBIND11_FINDPYTHON
        # EXAMPLE_VERSION_INFO shows you how to pass a value into the C++ code
        # from Python.
        cmake_args = [
            f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}{os.sep}",
            f"-DPYTHON_EXECUTABLE={sys.executable}",
            f"-DCMAKE_BUILD_TYPE={cfg}",  # not used on MSVC, but no harm
        ]
Xiaodong Ye's avatar
Xiaodong Ye committed
491
492
493
494
495

        if CUDA_HOME is not None:
            cmake_args += ["-DKTRANSFORMERS_USE_CUDA=ON"]
        elif MUSA_HOME is not None:
            cmake_args += ["-DKTRANSFORMERS_USE_MUSA=ON"]
Azure-Tang's avatar
Azure-Tang committed
496
497
        elif ROCM_HOME is not None:
            cmake_args += ["-DKTRANSFORMERS_USE_ROCM=ON"]
Xiaodong Ye's avatar
Xiaodong Ye committed
498
        else:
499
            raise ValueError("Unsupported backend: CUDA_HOME, MUSA_HOME, and ROCM_HOME are not set.")
500
501
        
        cmake_args = get_cmake_abi_args(cmake_args)
Azure-Tang's avatar
Azure-Tang committed
502
503
        # log cmake_args
        print("CMake args:", cmake_args)
504

chenxl's avatar
chenxl committed
505
506
        build_args = []
        if "CMAKE_ARGS" in os.environ:
chenxl's avatar
chenxl committed
507
508
            cmake_args += [
                item for item in os.environ["CMAKE_ARGS"].split(" ") if item]
Xiaodong Ye's avatar
Xiaodong Ye committed
509

510
511
512
513
514
515
516
517
        if CpuInstructInfo.CPU_INSTRUCT == CpuInstructInfo.FANCY:
            cpu_args = CpuInstructInfo.CMAKE_FANCY
        elif CpuInstructInfo.CPU_INSTRUCT == CpuInstructInfo.AVX512:
            cpu_args = CpuInstructInfo.CMAKE_AVX512
        elif CpuInstructInfo.CPU_INSTRUCT == CpuInstructInfo.AVX2:
            cpu_args = CpuInstructInfo.CMAKE_AVX2
        else:
            cpu_args = CpuInstructInfo.CMAKE_NATIVE
Xiaodong Ye's avatar
Xiaodong Ye committed
518

519
520
521
        cmake_args += [
            item for item in cpu_args.split(" ") if item
        ]
chenxl's avatar
chenxl committed
522
        # In this example, we pass in the version to C++. You might not need to.
chenxl's avatar
chenxl committed
523
524
        cmake_args += [
            f"-DEXAMPLE_VERSION_INFO={self.distribution.get_version()}"]
chenxl's avatar
chenxl committed
525
526
        if self.compiler.compiler_type != "msvc":
            if not cmake_generator or cmake_generator == "Ninja":
527
528
529
530
531
532
533
534
535
536
537
                pass
                # try:
                #     import ninja

                #     ninja_executable_path = Path(ninja.BIN_DIR) / "ninja"
                #     cmake_args += [
                #         "-GNinja",
                #         f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ninja_executable_path}",
                #     ]
                # except ImportError:
                #     pass
chenxl's avatar
chenxl committed
538
539
540

        else:
            # Single config generators are handled "normally"
chenxl's avatar
chenxl committed
541
542
            single_config = any(
                x in cmake_generator for x in {"NMake", "Ninja"})
chenxl's avatar
chenxl committed
543
544
545

            # CMake allows an arch-in-generator style for backward compatibility
            contains_arch = any(x in cmake_generator for x in {"ARM", "Win64"})
546
            if not single_config and not contains_arch and cmake_generator:
chenxl's avatar
chenxl committed
547
548
549
550
551
552
553
554
555
556
557
558
559
                cmake_args += ["-A", PLAT_TO_CMAKE[self.plat_name]]

            # Multi-config generators have a different way to specify configs
            if not single_config:
                cmake_args += [
                    f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{cfg.upper()}={extdir}"
                ]
                build_args += ["--config", cfg]

        if sys.platform.startswith("darwin"):
            # Cross-compile support for macOS - respect ARCHFLAGS if set
            archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", ""))
            if archs:
chenxl's avatar
chenxl committed
560
561
                cmake_args += [
                    "-DCMAKE_OSX_ARCHITECTURES={}".format(";".join(archs))]
chenxl's avatar
chenxl committed
562
563

        if "CMAKE_BUILD_PARALLEL_LEVEL" not in os.environ:
miaooo0000OOOO's avatar
miaooo0000OOOO committed
564
565
566
            cpu_count = os.cpu_count()
            if cpu_count is None:
                cpu_count = 1
chenxl's avatar
chenxl committed
567
            if hasattr(self, "parallel") and self.parallel:
miaooo0000OOOO's avatar
miaooo0000OOOO committed
568
569
570
                build_args += [f"--parallel={self.parallel}"]
            else:
                build_args += [f"--parallel={cpu_count}"]
liam's avatar
liam committed
571
        print("CMake args:", cmake_args)
chenxl's avatar
chenxl committed
572
        build_temp = Path(ext.sourcedir) / "build"
573
574
        print("build_temp:", build_temp)

chenxl's avatar
chenxl committed
575
576
        if not build_temp.exists():
            build_temp.mkdir(parents=True)
577
578
        run_command_with_live_tail(ext.name,
            ["cmake", ext.sourcedir, *cmake_args], cwd=build_temp
chenxl's avatar
chenxl committed
579
        )
580
581
        run_command_with_live_tail(ext.name,
            ["cmake", "--build", build_temp, "--verbose", *build_args], cwd=build_temp
chenxl's avatar
chenxl committed
582
583
        )

Azure-Tang's avatar
Azure-Tang committed
584
if CUDA_HOME is not None or ROCM_HOME is not None:
Xiaodong Ye's avatar
Xiaodong Ye committed
585
    ops_module = CUDAExtension('KTransformersOps', [
586
587
588
        'csrc/ktransformers_ext/cuda/custom_gguf/dequant.cu',
        'csrc/ktransformers_ext/cuda/binding.cpp',
        'csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cu'
Xiaodong Ye's avatar
Xiaodong Ye committed
589
590
591
592
593
    ],
    extra_compile_args={
            'cxx': ['-O3', '-DKTRANSFORMERS_USE_CUDA'],
            'nvcc': [
                '-O3',
Azure-Tang's avatar
Azure-Tang committed
594
                # '--use_fast_math',
Xiaodong Ye's avatar
Xiaodong Ye committed
595
596
597
598
599
600
                '-Xcompiler', '-fPIC',
                '-DKTRANSFORMERS_USE_CUDA',
            ]
        }
    )
elif MUSA_HOME is not None:
601
    SimplePorting(cuda_dir_path="csrc/ktransformers_ext/cuda", mapping_rule={
Xiaodong Ye's avatar
Xiaodong Ye committed
602
603
604
605
        # Common rules
        "at::cuda": "at::musa",
        "#include <ATen/cuda/CUDAContext.h>": "#include \"torch_musa/csrc/aten/musa/MUSAContext.h\"",
        "#include <c10/cuda/CUDAGuard.h>": "#include \"torch_musa/csrc/core/MUSAGuard.h\"",
Xiaodong Ye's avatar
Xiaodong Ye committed
606
        "nv_bfloat16": "mt_bfloat16",
Xiaodong Ye's avatar
Xiaodong Ye committed
607
608
        }).run()
    ops_module = MUSAExtension('KTransformersOps', [
609
610
        'csrc/ktransformers_ext/cuda_musa/custom_gguf/dequant.mu',
        'csrc/ktransformers_ext/cuda_musa/binding.cpp',
Xiaodong Ye's avatar
Xiaodong Ye committed
611
        # TODO: Add Marlin support for MUSA.
612
        # 'csrc/ktransformers_ext/cuda_musa/gptq_marlin/gptq_marlin.mu'
Xiaodong Ye's avatar
Xiaodong Ye committed
613
614
615
616
617
618
619
620
621
622
623
624
    ],
    extra_compile_args={
            'cxx': ['force_mcc'],
            'mcc': [
                '-O3',
                '-DKTRANSFORMERS_USE_MUSA',
                '-DTHRUST_IGNORE_CUB_VERSION_CHECK',
            ]
        }
    )
else:
    raise ValueError("Unsupported backend: CUDA_HOME and MUSA_HOME are not set.")
chenxl's avatar
chenxl committed
625

626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
ext_modules = [
    CMakeExtension("cpuinfer_ext", os.fspath(Path("").resolve() / "csrc" / "ktransformers_ext")),
    ops_module,
    CUDAExtension(
        'vLLMMarlin', [
            'csrc/custom_marlin/binding.cpp',
            'csrc/custom_marlin/gptq_marlin/gptq_marlin.cu',
            'csrc/custom_marlin/gptq_marlin/gptq_marlin_repack.cu',
        ],
        extra_compile_args={
            'cxx': ['-O3'],
            'nvcc': ['-O3', '-Xcompiler', '-fPIC'],
        },
    )
]
if with_balance:
    print("using balance_serve")
    ext_modules.append(
        CMakeExtension("balance_serve", os.fspath(Path("").resolve()/ "csrc"/ "balance_serve"))
    )

chenxl's avatar
chenxl committed
647
setup(
Azure-Tang's avatar
Azure-Tang committed
648
    name=VersionInfo.PACKAGE_NAME,
chenxl's avatar
chenxl committed
649
    version=VersionInfo().get_package_version(),
chenxl's avatar
chenxl committed
650
    cmdclass={"bdist_wheel":BuildWheelsCommand ,"build_ext": CMakeBuild},
651
    ext_modules=ext_modules
652
)