[Feature] Qwen 模型 Triton SwiGLU 激活函数优化 (2.58x 加速)

src/mindnlp/triton/__init__.py

@@ -1,57 +1,128 @@
 """triton adapter for mindspore"""
+import os as _os
+import sys as _sys
 from functools import lru_cache
-import mindspore
-from mindtorch import ops
-from mindnlp.utils import is_triton_available
 
-if is_triton_available():
-    from triton.backends.driver import DriverBase  # pylint: disable=import-error
-    from triton.backends.nvidia.driver import CudaUtils, CudaLauncher  # pylint: disable=import-error
-    from triton.backends.compiler import GPUTarget  # pylint: disable=import-error
+# Set TRITON_BACKEND to mindspore to avoid torch_npu._C dependency issues
+# when MindTorchFinder proxies torch_npu imports
+if _os.environ.get('TRITON_BACKEND') not in ('torch_npu', 'mindspore'):
+    _os.environ['TRITON_BACKEND'] = 'mindspore'
 
-class MSDriver(DriverBase):
 
-    def __init__(self):
-        self.utils = CudaUtils()  # TODO: make static
-        self.launcher_cls = CudaLauncher
-        super().__init__()
+def _get_mindspore():
+    """Lazy import mindspore to avoid import-time hanging."""
+    import mindspore
+    return mindspore
 
-    def get_current_device(self):
-        return 0
 
-    def set_current_device(self):
+def _get_mindtorch_ops():
+    """Lazy import mindtorch.ops."""
+    from mindtorch import ops
+    return ops
+
+
+def _register_ms_driver():
+    """Register MSDriver for Triton if triton is available."""
+    try:
+        from mindnlp.utils import is_triton_available
+        if not is_triton_available():
+            return
+    except ImportError:
+        return
+
+    try:
+        from triton.backends.driver import DriverBase
+        from triton.backends.nvidia.driver import CudaUtils, CudaLauncher
+        from triton.backends.compiler import GPUTarget
+
+        class MSDriver(DriverBase):
+
+            def __init__(self):
+                self.utils = CudaUtils()
+                self.launcher_cls = CudaLauncher
+                super().__init__()
+
+            def get_current_device(self):
+                return 0
+
+            def set_current_device(self):
+                pass
+
+            @lru_cache
+            def get_current_stream(self, device=None):
+                ms = _get_mindspore()
+                return ms.hal.current_stream().id
+
+            @lru_cache
+            def get_device_capability(self, device=0):
+                ms = _get_mindspore()
+                return ms.hal.get_device_capability(0)
+
+            @lru_cache
+            def get_current_target(self):
+                device = self.get_current_device()
+                capability = self.get_device_capability(device)
+                capability = capability[0] * 10 + capability[1]
+                warp_size = 32
+                return GPUTarget("cuda", capability, warp_size)
+
+            def get_device_interface(self):
+                ms = _get_mindspore()
+                return ms.hal
+
+            @staticmethod
+            def is_active():
+                return True
+
+            def get_benchmarker(self):
+                from triton.testing import do_bench
+                return do_bench
+
+            def get_empty_cache_for_benchmark(self):
+                cache_size = 256 * 1024 * 1024
+                ops = _get_mindtorch_ops()
+                ms = _get_mindspore()
+                return ops.empty(int(cache_size // 4), dtype=ms.int32, device='GPU')
+
+    except (ImportError, AttributeError):
         pass
 
-    @lru_cache
-    def get_current_stream(self, device=None):
-        return mindspore.hal.current_stream().id
-
-    @lru_cache
-    def get_device_capability(self, device=0):
-        return mindspore.hal.get_device_capability(0)
-
-    @lru_cache
-    def get_current_target(self):
-        device = self.get_current_device()
-        capability = self.get_device_capability(device)
-        capability = capability[0] * 10 + capability[1]
-        warp_size = 32
-        return GPUTarget("cuda", capability, warp_size)
-
-    def get_device_interface(self):
-        return mindspore.hal
-
-    @staticmethod
-    def is_active():
-        return True
-
-    def get_benchmarker(self):
-        from triton.testing import do_bench  # pylint: disable=import-error
-        return do_bench
-
-    def get_empty_cache_for_benchmark(self):
-        # We maintain a buffer of 256 MB that we clear
-        # before each kernel call to make sure that the L2 cache
-        # doesn't contain any input data before the run
-        cache_size = 256 * 1024 * 1024
-        return ops.empty(int(cache_size // 4), dtype=mindspore.int32, device='GPU')
+
+_register_ms_driver()
+
+__all__ = [
+    "TritonGELU",
+    "TritonSwiGLU",
+    "triton_gelu",
+    "triton_swiglu",
+    "native_gelu",
+    "native_swiglu",
+    "get_available_backend",
+    "MSGELU",
+    "MSSwiGLU",
+    "MSTritonGELU",
+    "MSTritonSwiGLU",
+    "gelu",
+    "swiglu",
+    "get_ms_activation",
+]
+
+def __getattr__(name):
+    if name in ("TritonGELU", "TritonSwiGLU", "triton_gelu", "triton_swiglu", "native_gelu", "native_swiglu"):
+        from mindnlp.triton.kernels.activations import (
+            TritonGELU, TritonSwiGLU, triton_gelu, triton_swiglu, native_gelu, native_swiglu,
+        )
+        globals().update(locals())
+        return locals()[name]
+    if name in ("MSGELU", "MSSwiGLU", "MSTritonGELU", "MSTritonSwiGLU", "gelu", "swiglu", "get_ms_activation"):
+        from mindnlp.triton.kernels.mindspore_adapter import (
+            MSGELU, MSSwiGLU, TritonGELU as MSTritonGELU, TritonSwiGLU as MSTritonSwiGLU,
+            gelu, swiglu, get_ms_activation,
+        )
+        globals().update(locals())
+        return locals()[name]
+    if name == "get_available_backend":
+        from mindnlp.triton.backends.detect import get_available_backend
+        globals()["get_available_backend"] = get_available_backend
+        return get_available_backend
+    raise AttributeError(f"module 'mindnlp.triton' has no attribute '{name}'")

src/mindnlp/triton/backends/__init__.py

@@ -0,0 +1,21 @@
+"""
+Triton Backend Implementations.
+"""
+
+from mindnlp.triton.backends.detect import (
+    get_available_backend,
+    is_triton_available,
+    BackendType,
+)
+from mindnlp.triton.backends.ascend import (
+    get_ascend_backend,
+    is_ascend_available,
+)
+
+__all__ = [
+    "BackendType",
+    "get_available_backend",
+    "is_triton_available",
+    "get_ascend_backend",
+    "is_ascend_available",
+]

src/mindnlp/triton/backends/ascend.py

@@ -0,0 +1,102 @@
+"""
+Ascend NPU backend implementation.
+"""
+
+from typing import Optional, List
+import torch
+
+from mindnlp.triton.backends.detect import BackendType, is_triton_available
+
+
+class AscendBackend:
+    """Ascend NPU backend for Triton kernels."""
+
+    def __init__(self):
+        self._device = "npu"
+        self._initialized = False
+
+    def initialize(self) -> bool:
+        """Initialize the Ascend backend.
+
+        Returns:
+            True if initialization succeeded, False otherwise
+        """
+        if self._initialized:
+            return True
+
+        try:
+            if not torch.npu.is_available():
+                return False
+
+            import triton
+            self._initialized = True
+            return True
+        except Exception:
+            return False
+
+    @property
+    def device(self) -> str:
+        """Return the device string for this backend."""
+        return self._device
+
+    @property
+    def device_count(self) -> int:
+        """Return the number of available devices."""
+        return torch.npu.device_count()
+
+    def synchronize(self, device_id: Optional[int] = None):
+        """Synchronize the device.
+
+        Args:
+            device_id: Optional device ID to synchronize
+        """
+        if device_id is not None:
+            with torch.npu.device(f"npu:{device_id}"):
+                torch.npu.synchronize()
+        else:
+            torch.npu.synchronize()
+
+    def set_device(self, device_id: int):
+        """Set the current device.
+
+        Args:
+            device_id: Device ID to set
+        """
+        torch.npu.set_device(device_id)
+
+
+_ascend_backend_instance: Optional[AscendBackend] = None
+
+
+def get_ascend_backend() -> AscendBackend:
+    """Get the global Ascend backend instance.
+
+    Returns:
+        AscendBackend instance
+    """
+    global _ascend_backend_instance
+    if _ascend_backend_instance is None:
+        _ascend_backend_instance = AscendBackend()
+    return _ascend_backend_instance
+
+
+def is_ascend_available() -> bool:
+    """Check if Ascend backend is available.
+
+    Returns:
+        True if Ascend NPU is available, False otherwise
+    """
+    return torch.npu.is_available() and is_triton_available()
+
+
+def is_triton_available_for_ascend() -> bool:
+    """Check if Triton is available and can target Ascend.
+
+    Returns:
+        True if Triton can target Ascend, False otherwise
+    """
+    try:
+        import triton
+        return hasattr(triton, 'runtime')
+    except Exception:
+        return False

src/mindnlp/triton/backends/detect.py

@@ -0,0 +1,89 @@
+"""
+Backend detection and selection.
+"""
+
+import os
+from enum import Enum
+from typing import Optional
+
+
+class BackendType(Enum):
+    """Available Triton backends."""
+    ASCEND = "ascend"
+    NVIDIA = "nvidia"
+    CPU = "cpu"
+    NONE = "none"
+
+
+def detect_npu() -> bool:
+    """Detect if Ascend NPU is available."""
+    try:
+        import torch
+        return torch.npu.is_available()
+    except Exception:
+        return False
+
+
+def detect_cuda() -> bool:
+    """Detect if NVIDIA GPU is available."""
+    try:
+        import torch
+        return torch.cuda.is_available()
+    except Exception:
+        return False
+
+
+def get_available_backend() -> BackendType:
+    """Detect and return the available Triton backend.
+
+    Returns:
+        BackendType enum value indicating which backend is available
+    """
+    env_backend = os.environ.get("MINNLP_TRITON_BACKEND", "").lower()
+
+    if env_backend == "ascend":
+        if detect_npu():
+            return BackendType.ASCEND
+    elif env_backend == "nvidia":
+        if detect_cuda():
+            return BackendType.NVIDIA
+    elif env_backend == "cpu":
+        return BackendType.CPU
+
+    if detect_npu():
+        return BackendType.ASCEND
+    if detect_cuda():
+        return BackendType.NVIDIA
+
+    return BackendType.NONE
+
+
+def is_triton_available() -> bool:
+    """Check if Triton is available."""
+    try:
+        import triton
+        return True
+    except ImportError:
+        return False
+
+
+def is_backend_available(backend: BackendType) -> bool:
+    """Check if a specific backend is available.
+
+    Args:
+        backend: Backend type to check
+
+    Returns:
+        True if the backend is available, False otherwise
+    """
+    if backend == BackendType.NONE:
+        return False
+
+    if backend == BackendType.ASCEND:
+        return detect_npu() and is_triton_available()
+    elif backend == BackendType.NVIDIA:
+        return detect_cuda() and is_triton_available()
+    elif backend == BackendType.CPU:
+        return is_triton_available()
+
+    return False