deepspeed

1
# Copyright (c) Microsoft Corporation.
2
# SPDX-License-Identifier: Apache-2.0
3

4
# DeepSpeed Team
5

6
from .builder import CUDAOpBuilder
7

8
import sys
9

10

11
class FusedLionBuilder(CUDAOpBuilder):
12
    BUILD_VAR = "DS_BUILD_FUSED_LION"
13
    NAME = "fused_lion"
14

15
    def __init__(self):
16
        super().__init__(name=self.NAME)
17

18
    def absolute_name(self):
19
        return f'deepspeed.ops.lion.{self.NAME}_op'
20

21
    def sources(self):
22
        return ['csrc/lion/fused_lion_frontend.cpp', 'csrc/lion/multi_tensor_lion.cu']
23

24
    def include_paths(self):
25
        return ['csrc/includes', 'csrc/lion']
26

27
    def cxx_args(self):
28
        args = super().cxx_args()
29
        return args + self.version_dependent_macros()
30

31
    def nvcc_args(self):
32
        nvcc_flags = ['-O3'] + self.version_dependent_macros()
33
        if not self.is_rocm_pytorch():
34
            nvcc_flags.extend(
35
                ['-allow-unsupported-compiler' if sys.platform == "win32" else '', '-lineinfo', '--use_fast_math'] +
36
                self.compute_capability_args())
37
        return nvcc_flags
38