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

9
class TransformerBuilder(CUDAOpBuilder):
10
    BUILD_VAR = "DS_BUILD_TRANSFORMER"
11
    NAME = "transformer"
12

13
    def __init__(self, name=None):
14
        name = self.NAME if name is None else name
15
        super().__init__(name=name)
16

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

20
    def extra_ldflags(self):
21
        if not self.is_rocm_pytorch():
22
            return ['-lcurand']
23
        else:
24
            return []
25

26
    def sources(self):
27
        return [
28
            'csrc/transformer/ds_transformer_cuda.cpp', 'csrc/transformer/cublas_wrappers.cu',
29
            'csrc/transformer/transform_kernels.cu', 'csrc/transformer/gelu_kernels.cu',
30
            'csrc/transformer/dropout_kernels.cu', 'csrc/transformer/normalize_kernels.cu',
31
            'csrc/transformer/softmax_kernels.cu', 'csrc/transformer/general_kernels.cu'
32
        ]
33

34
    def include_paths(self):
35
        includes = ['csrc/includes']
36
        return includes
37