sparse_attn.py

"""
Copyright 2020 The Microsoft DeepSpeed Team
"""
import torch
import warnings
from .builder import OpBuilder


class SparseAttnBuilder(OpBuilder):
    BUILD_VAR = "DS_BUILD_SPARSE_ATTN"
    NAME = "sparse_attn"

    def __init__(self):
        super().__init__(name=self.NAME)

    def absolute_name(self):
        return f'deepspeed.ops.sparse_attention.{self.NAME}_op'

    def sources(self, is_rocm_pytorch):
        if is_rocm_pytorch:
            return ['csrc/sparse_attention/hip/utils.cpp']
        else:
            return ['csrc/sparse_attention/utils.cpp']

    def cxx_args(self):
        return ['-O2', '-fopenmp']
    
    def is_compatible(self):
        # Check to see if llvm and cmake are installed since they are dependencies
        #required_commands = ['llvm-config|llvm-config-9', 'cmake']
        required_commands = ['cmake']
        command_status = list(map(self.command_exists, required_commands))
        deps_compatible = all(command_status)

        # torch-cpu will not have a cuda version
        if torch.version.hip is None:
            cuda_compatible = False
            self.warning(f"{self.NAME} cuda is not available from torch")
        #else:
        #    major, minor = torch.version.cuda.split('.')[:2]
        #    cuda_compatible = int(major) == 10 and int(minor) >= 1
        #    if not cuda_compatible:
        #        self.warning(
        #            f"{self.NAME} requires CUDA version 10.1+, does not currently support >=11 or <10.1"
        #        )

        TORCH_MAJOR = int(torch.__version__.split('.')[0])
        TORCH_MINOR = int(torch.__version__.split('.')[1])
        torch_compatible = TORCH_MAJOR == 1 and TORCH_MINOR >= 5
        if not torch_compatible:
            self.warning(
                f'{self.NAME} requires a torch version >= 1.5 but detected {TORCH_MAJOR}.{TORCH_MINOR}'
            )

        return super().is_compatible(
        ) and deps_compatible and torch_compatible #and cuda_compatible