[NPU] Add capability-based Gemma RMSNorm dispatch

python/sglang/srt/hardware_backend/npu/device_capabilities.py

@@ -0,0 +1,65 @@
+"""Semantic feature support for Ascend NPU device families."""
+
+import logging
+from enum import Enum
+from functools import lru_cache
+
+logger = logging.getLogger(__name__)
+
+
+class NPUDeviceFamily(str, Enum):
+    ASCEND_910B = "ascend_910b"
+    ASCEND_910C = "ascend_910c"
+    ASCEND_950 = "ascend_950"
+    UNKNOWN = "unknown"
+
+
+class NPUFeature(str, Enum):
+    NATIVE_GEMMA_RMS_NORM = "native_gemma_rms_norm"
+    TRITON_GEMMA_RMS_NORM = "triton_gemma_rms_norm"
+
+
+_DEVICE_FAMILY_RANGES = (
+    (220, 225, NPUDeviceFamily.ASCEND_910B),
+    (250, 255, NPUDeviceFamily.ASCEND_910C),
+    (260, 260, NPUDeviceFamily.ASCEND_950),
+)
+
+_FEATURES_BY_DEVICE_FAMILY = {
+    NPUDeviceFamily.ASCEND_910B: frozenset({NPUFeature.NATIVE_GEMMA_RMS_NORM}),
+    NPUDeviceFamily.ASCEND_910C: frozenset({NPUFeature.NATIVE_GEMMA_RMS_NORM}),
+    NPUDeviceFamily.ASCEND_950: frozenset({NPUFeature.TRITON_GEMMA_RMS_NORM}),
+    NPUDeviceFamily.UNKNOWN: frozenset(),
+}
+
+
+@lru_cache(maxsize=1)
+def get_npu_device_family() -> NPUDeviceFamily:
+    """Resolve and cache the current process-wide Ascend device family."""
+    try:
+        import torch_npu
+
+        soc_version = int(torch_npu.npu.get_soc_version())
+    except (ImportError, AttributeError, RuntimeError, TypeError, ValueError) as e:
+        logger.warning(
+            "Unable to resolve the Ascend SoC version; optional NPU features "
+            "will use safe fallbacks: %s",
+            e,
+        )
+        return NPUDeviceFamily.UNKNOWN
+
+    for lower, upper, family in _DEVICE_FAMILY_RANGES:
+        if lower <= soc_version <= upper:
+            return family
+
+    logger.warning(
+        "Unknown Ascend SoC version %s; optional NPU features will use safe "
+        "fallbacks.",
+        soc_version,
+    )
+    return NPUDeviceFamily.UNKNOWN
+
+
+def supports_npu_feature(feature: NPUFeature) -> bool:
+    """Whether the current Ascend family supports a semantic NPU feature."""
+    return feature in _FEATURES_BY_DEVICE_FAMILY[get_npu_device_family()]

python/sglang/srt/hardware_backend/npu/device_operator.py

@@ -0,0 +1,103 @@
+"""Device-aware operator facade for Ascend NPU kernels."""
+
+from typing import TYPE_CHECKING, Tuple
+
+from sglang.srt.hardware_backend.npu.device_capabilities import (
+    NPUFeature,
+    supports_npu_feature,
+)
+
+if TYPE_CHECKING:
+    import torch
+
+
+def _triton_gemma_rms_norm(
+    input: "torch.Tensor",
+    weight: "torch.Tensor",
+    residual: "torch.Tensor | None",
+    eps: float,
+) -> Tuple["torch.Tensor", "torch.Tensor | None"]:
+    if input.numel() == 0:
+        residual_sum = None if residual is None else input + residual
+        return input.clone(), residual_sum
+
+    from sgl_kernel_npu.norm.add_rmsnorm_bias import add_gemma_rms_norm
+
+    original_shape = input.shape
+    input_2d = input.contiguous().view(-1, original_shape[-1])
+    residual_2d = None
+    if residual is not None:
+        residual_2d = residual.contiguous().view_as(input_2d)
+
+    norm_output, residual_sum = add_gemma_rms_norm(
+        input_2d,
+        weight.contiguous(),
+        residual_2d,
+        eps,
+    )
+    norm_output = norm_output.view(original_shape)
+    if residual is not None:
+        residual_sum = residual_sum.view(original_shape)
+    return norm_output, residual_sum
+
+
+class NPUDeviceOperator:
+    """Expose semantic operator APIs across Ascend device families."""
+
+    @staticmethod
+    def gemma_rms_norm(
+        input: "torch.Tensor",
+        weight: "torch.Tensor",
+        eps: float,
+    ) -> "torch.Tensor":
+        """Apply Gemma RMSNorm without exposing product checks to callers."""
+        if input.numel() == 0:
+            return input.clone()
+
+        input = input.contiguous()
+        weight = weight.contiguous()
+        if supports_npu_feature(NPUFeature.NATIVE_GEMMA_RMS_NORM):
+            import torch_npu
+
+            return torch_npu.npu_gemma_rms_norm(input, weight, eps)[0]
+
+        if supports_npu_feature(NPUFeature.TRITON_GEMMA_RMS_NORM):
+            return _triton_gemma_rms_norm(input, weight, None, eps)[0]
+
+        import torch_npu
+
+        return torch_npu.npu_rms_norm(input, 1.0 + weight, eps)[0]
+
+    @staticmethod
+    def add_gemma_rms_norm(
+        input: "torch.Tensor",
+        weight: "torch.Tensor",
+        residual: "torch.Tensor",
+        eps: float,
+    ) -> Tuple["torch.Tensor", "torch.Tensor"]:
+        """Add a residual and apply Gemma RMSNorm without mutating inputs."""
+        if input.numel() == 0:
+            return input.clone(), input + residual
+
+        input = input.contiguous()
+        weight = weight.contiguous()
+        residual = residual.contiguous()
+        if supports_npu_feature(NPUFeature.NATIVE_GEMMA_RMS_NORM):
+            import torch_npu
+
+            norm_output, _, residual_sum = torch_npu.npu_add_rms_norm(
+                residual, input, 1.0 + weight, eps
+            )
+            return norm_output, residual_sum
+
+        if supports_npu_feature(NPUFeature.TRITON_GEMMA_RMS_NORM):
+            norm_output, residual_sum = _triton_gemma_rms_norm(
+                input, weight, residual, eps
+            )
+            return norm_output, residual_sum
+
+        import torch_npu
+
+        residual_sum = input + residual
+        norm_output = torch_npu.npu_rms_norm(residual_sum, 1.0 + weight, eps)[0]
+        return norm_output, residual_sum

python/sglang/srt/layers/layernorm.py

@@ -26,6 +26,7 @@
     rms_norm_batch_invariant,
 )
 from sglang.srt.environ import envs
+from sglang.srt.hardware_backend.npu.device_operator import NPUDeviceOperator
 from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.model_executor.cuda_graph_config import (
     Backend,
@@ -159,7 +160,6 @@ def is_supported_rmsnorm_hf_hidden_size(d: int) -> bool:
 
 if _is_npu:
     import torch_npu
-    from sgl_kernel_npu.norm.add_rmsnorm_bias import add_gemma_rms_norm
 
 
 @lru_cache(maxsize=1)
@@ -995,13 +995,11 @@ def forward_npu(
         if residual is not None:
             if post_residual_addition is not None:
                 residual = residual + post_residual_addition
-            norm_out, residual = add_gemma_rms_norm(
+            return NPUDeviceOperator.add_gemma_rms_norm(
                 x, self.weight, residual, self.variance_epsilon
             )
-            return norm_out, residual
 
-        x, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.variance_epsilon)
-        return x
+        return NPUDeviceOperator.gemma_rms_norm(x, self.weight, self.variance_epsilon)
 
     def forward_xpu(
         self,
@@ -1093,10 +1091,13 @@ def forward_hip(self, x, residual: Optional[torch.Tensor] = None):
         return self.forward_native(x, residual)
 
     def forward_npu(self, x, residual: Optional[torch.Tensor] = None):
-        if residual is not None:
+        if envs.SGLANG_NPU_FORWARD_NATIVE_GEMMA_RMS_NORM.get():
             return self.forward_native(x, residual)
-        output, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.eps)
-        return output
+        if residual is not None:
+            return NPUDeviceOperator.add_gemma_rms_norm(
+                x, self.weight, residual, self.eps
+            )
+        return NPUDeviceOperator.gemma_rms_norm(x, self.weight, self.eps)
 
     def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.eps}"