deepspeedai--deepspeed/op_builder/evoformer_attn.py

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

# DeepSpeed Team

from .builder import CUDAOpBuilder, installed_cuda_version
import importlib
import os
from pathlib import Path
import sys


class EvoformerAttnBuilder(CUDAOpBuilder):
    BUILD_VAR = "DS_BUILD_EVOFORMER_ATTN"
    NAME = "evoformer_attn"
    CUTLASS_IGNORE = "DS_IGNORE_CUTLASS_DETECTION"
    CUTLASS_PYTHON_BINDINGS = "DS_USE_CUTLASS_PYTHON_BINDINGS"

    def __init__(self, name=None):
        name = self.NAME if name is None else name
        super().__init__(name=name)
        self.cutlass_path = os.environ.get("CUTLASS_PATH")
        self._resolved_cutlass_path = None

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

    def extra_ldflags(self):
        if not self.is_rocm_pytorch():
            return ["-lcurand"]
        else:
            return []

    def sources(self):
        src_dir = "csrc/deepspeed4science/evoformer_attn"
        return [f"{src_dir}/attention.cpp", f"{src_dir}/attention_back.cu", f"{src_dir}/attention_cu.cu"]

    def nvcc_args(self):
        if os.environ.get("DS_EVOFORMER_GPU_ARCH"):
            self.warning("DS_EVOFORMER_GPU_ARCH is deprecated and ignored for Evoformer builds. "
                         "Use TORCH_CUDA_ARCH_LIST to control build targets.")
        return super().nvcc_args()

    def filter_ccs(self, ccs):
        """Keep only Tensor Core capable targets (>= 7.0)."""
        retained = []
        pruned = []
        for cc in [cc.split('.') for cc in ccs]:
            if int(cc[0]) >= 7:
                retained.append(cc)
            else:
                pruned.append(cc)
        if pruned:
            self.warning(f"Evoformer: excluding targets below SM 7.0: {pruned}. Tensor Core required.")
        return retained

    def is_compatible(self, verbose=False):
        try:
            import torch
        except ImportError:
            if verbose:
                self.warning("Please install torch if trying to pre-compile kernels")
            return False

        if self.cutlass_path != self.CUTLASS_IGNORE:
            try:
                self.include_paths()
            except (RuntimeError, ImportError) as exc:
                if verbose:
                    self.warning(str(exc))
                return False
            # Check version in case it is a CUTLASS_PATH points to a CUTLASS checkout
            if self._resolved_cutlass_path is not None:
                changelog_path = self._resolved_cutlass_path / "CHANGELOG.md"
            else:
                changelog_path = None
            if changelog_path is not None and changelog_path.exists():
                with open(changelog_path, "r") as f:
                    if "3.1.0" not in f.read():
                        if verbose:
                            self.warning("Please use CUTLASS version >= 3.1.0")
                        return False

        # Check CUDA and GPU capabilities
        cuda_okay = True
        if not os.environ.get("DS_IGNORE_CUDA_DETECTION"):
            if not self.is_rocm_pytorch() and torch.cuda.is_available():  #ignore-cuda
                sys_cuda_major, _ = installed_cuda_version()
                torch_cuda_major = int(torch.version.cuda.split(".")[0])
                cuda_capability = self.cuda_capability_major()
                if cuda_capability is not None and cuda_capability < 7:
                    if verbose:
                        self.warning("Please use a GPU with compute capability >= 7.0")
                    cuda_okay = False
                if torch_cuda_major < 11 or sys_cuda_major < 11:
                    if verbose:
                        self.warning("Please use CUDA 11+")
                    cuda_okay = False
        return super().is_compatible(verbose) and cuda_okay

    @staticmethod
    def _repo_root():
        return Path(__file__).resolve().parents[1]

    @staticmethod
    def _dedupe_paths(paths):
        deduped = []
        seen = set()
        for path in paths:
            path = Path(path).expanduser()
            key = str(path)
            if key not in seen:
                seen.add(key)
                deduped.append(path)
        return deduped

    @staticmethod
    def _env_paths(*names):
        paths = []
        for name in names:
            value = os.environ.get(name)
            if not value:
                continue
            paths.extend(Path(path) for path in value.split(os.pathsep) if path)
        return paths

    @staticmethod
    def _python_package_cutlass_paths():
        try:
            cutlass_library = importlib.import_module("cutlass_library")
        except ImportError:
            return []

        candidates = []
        source_path = getattr(cutlass_library, "source_path", None)
        if source_path is not None:
            candidates.append(Path(source_path))

        package_file = getattr(cutlass_library, "__file__", None)
        if package_file is not None:
            package_dir = Path(package_file).resolve().parent
            candidates.extend([package_dir / "source", package_dir.parent, package_dir])
        return candidates

    def _candidate_cutlass_paths(self):
        if self.cutlass_path == self.CUTLASS_PYTHON_BINDINGS:
            candidates = self._python_package_cutlass_paths()
            if candidates:
                return candidates
            self.warning("Please pip install nvidia-cutlass")
            raise ImportError("Unable to locate CUTLASS from the nvidia-cutlass Python package")

        if self.cutlass_path:
            return [Path(self.cutlass_path)]

        repo_root = self._repo_root()
        python_prefixes = self._dedupe_paths([Path(sys.prefix), Path(sys.exec_prefix), Path(sys.base_prefix)])
        prefix_paths = self._env_paths("CUTLASS_ROOT", "CUTLASS_HOME", "CONDA_PREFIX", "VIRTUAL_ENV",
                                       "CMAKE_PREFIX_PATH", "CUDA_HOME", "CUDA_PATH")
        include_paths = self._env_paths("CPATH", "CPLUS_INCLUDE_PATH", "C_INCLUDE_PATH")

        return self._dedupe_paths([
            *self._python_package_cutlass_paths(),
            *prefix_paths,
            *python_prefixes,
            *include_paths,
            Path.cwd() / "cutlass",
            repo_root / "cutlass",
            repo_root.parent / "cutlass",
            Path("/usr/local/cutlass"),
            Path("/opt/cutlass"),
            Path("/usr/local"),
            Path("/usr"),
        ])

    @staticmethod
    def _cutlass_include_dirs(cutlass_path):
        cutlass_path = cutlass_path.expanduser().resolve()
        if not cutlass_path.is_dir():
            return []

        if (cutlass_path / "include" / "cutlass" / "cutlass.h").is_file():
            include_root = cutlass_path / "include"
            util_include = cutlass_path / "tools" / "util" / "include"
        elif (cutlass_path / "cutlass" / "cutlass.h").is_file():
            include_root = cutlass_path
            util_include = cutlass_path.parent / "tools" / "util" / "include"
        else:
            return []

        include_dirs = [include_root]
        if util_include.is_dir():
            include_dirs.append(util_include)
        return [str(include_dir) for include_dir in include_dirs]

    def include_paths(self):
        # Assume the user knows best and CUTLASS location is already setup externally
        if self.cutlass_path == self.CUTLASS_IGNORE:
            return []

        for cutlass_path in self._candidate_cutlass_paths():
            include_dirs = self._cutlass_include_dirs(cutlass_path)
            if include_dirs:
                self._resolved_cutlass_path = cutlass_path.expanduser().resolve()
                return include_dirs

        if self.cutlass_path:
            raise RuntimeError(f"CUTLASS_PATH {self.cutlass_path} does not contain CUTLASS headers")

        raise RuntimeError("Unable to locate CUTLASS. Install nvidia-cutlass, clone CUTLASS next to DeepSpeed, "
                           "or set CUTLASS_PATH to the CUTLASS checkout.")