[Kernel] Fill non-CUDA coverage: HIP (aiter/rocm-triton) + Ascend NPU backends (RFC #29630) (#31307)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

python/sglang/kernels/fused_op.py

@@ -70,6 +70,7 @@ BACKEND_METHODS: Dict[KernelBackend, str] = {
     KernelBackend.FLASHINFER: "forward_flashinfer",
     KernelBackend.DEEPGEMM: "forward_deepgemm",
     KernelBackend.AITER: "forward_aiter",
+    KernelBackend.TORCH_NPU: "forward_npu",
 }
 
 # best -> fallback. ``torch_compile`` is deliberately absent: auto-selection
@@ -82,6 +83,7 @@ DEFAULT_PRIORITY: Tuple[KernelBackend, ...] = (
     KernelBackend.DEEPGEMM,
     KernelBackend.CUTE_DSL,
     KernelBackend.AITER,
+    KernelBackend.TORCH_NPU,
     KernelBackend.TRITON,
     KernelBackend.TORCH,
 )
@@ -269,6 +271,9 @@ class BaseFusedOp(ABC):
     def forward_aiter(self, *args, **kwargs):
         raise NotImplementedError(f"{self.op}: no aiter backend")
 
+    def forward_npu(self, *args, **kwargs):
+        raise NotImplementedError(f"{self.op}: no npu backend")
+
     # --- selection ---
 
     def available_backends(self) -> List[KernelBackend]:

python/sglang/kernels/ops/activation/__init__.py

@@ -180,11 +180,98 @@ class GeluTanhAndMulOp(_GatedActivationOp):
         return F.gelu(gate, approximate="tanh")
 
 
+class ReLU2Op(BaseFusedOp):
+    """``out = relu(input) ** 2`` (single-input, not gated).
+
+    The real kernel is the CUDA JIT path (``sglang.jit_kernel.activation.relu2``,
+    used in production on CUDA); elsewhere the torch reference runs.
+    """
+
+    op = "activation.relu2"
+    priority = (KernelBackend.JIT, KernelBackend.TORCH)
+    capabilities = {KernelBackend.JIT: _CUDA}
+    format_signature = FormatSignature(
+        supported_dtypes=_ACT_DTYPES,
+        description="relu(x) ** 2; returns tensor",
+    )
+    descriptions = {
+        KernelBackend.JIT: "relu(x)**2 (sglang.jit_kernel).",
+        KernelBackend.TORCH: "relu(x)**2 (pure-torch reference).",
+    }
+
+    def forward_native(
+        self, input: torch.Tensor, out: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        import torch.nn.functional as F
+
+        x = F.relu(input)
+        result = x * x
+        if out is None:
+            return result
+        out.copy_(result)
+        return out
+
+    def forward_jit(
+        self, input: torch.Tensor, out: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        from sglang.jit_kernel.activation import relu2
+
+        result = relu2(input)
+        if out is None:
+            return result
+        out.copy_(result)
+        return out
+
+
+class QuickGELUOp(BaseFusedOp):
+    """``out = input * sigmoid(1.702 * input)`` (single-input, not gated).
+
+    Only ROCm has a native kernel (``sgl_kernel.gelu_quick``, AOT on HIP); CUDA
+    uses the torch reference in production, so no CUDA backend is registered.
+    """
+
+    op = "activation.gelu_quick"
+    priority = (KernelBackend.AOT, KernelBackend.TORCH)
+    capabilities = {KernelBackend.AOT: _HIP}
+    format_signature = FormatSignature(
+        supported_dtypes=_ACT_DTYPES,
+        description="x * sigmoid(1.702 * x); returns tensor",
+    )
+    descriptions = {
+        KernelBackend.AOT: "quick GELU (sgl_kernel wheel, ROCm).",
+        KernelBackend.TORCH: "quick GELU (pure-torch reference).",
+    }
+
+    def forward_native(
+        self, input: torch.Tensor, out: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        import torch
+
+        result = input * torch.sigmoid(1.702 * input)
+        if out is None:
+            return result
+        out.copy_(result)
+        return out
+
+    def forward_aot(
+        self, input: torch.Tensor, out: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        import torch
+        from sgl_kernel import gelu_quick
+
+        if out is None:
+            out = torch.empty(input.shape, dtype=input.dtype, device=input.device)
+        gelu_quick(input, out)
+        return out
+
+
 _SILU_AND_MUL = register_fused_op(SiluAndMulOp(), __name__, "_SILU_AND_MUL")
 _GELU_AND_MUL = register_fused_op(GeluAndMulOp(), __name__, "_GELU_AND_MUL")
 _GELU_TANH_AND_MUL = register_fused_op(
     GeluTanhAndMulOp(), __name__, "_GELU_TANH_AND_MUL"
 )
+_RELU2 = register_fused_op(ReLU2Op(), __name__, "_RELU2")
+_GELU_QUICK = register_fused_op(QuickGELUOp(), __name__, "_GELU_QUICK")
 
 
 def silu_and_mul(
@@ -208,13 +295,27 @@ def gelu_tanh_and_mul(
     return _GELU_TANH_AND_MUL(input, out)
 
 
+def relu2(input: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """``out = relu(input) ** 2``."""
+    return _RELU2(input, out)
+
+
+def gelu_quick(input: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """``out = input * sigmoid(1.702 * input)``."""
+    return _GELU_QUICK(input, out)
+
+
 __all__ = [
     "SiluAndMulOp",
     "GeluAndMulOp",
     "GeluTanhAndMulOp",
+    "ReLU2Op",
+    "QuickGELUOp",
     "silu_and_mul",
     "gelu_and_mul",
     "gelu_tanh_and_mul",
+    "relu2",
+    "gelu_quick",
 ]
 
 

python/sglang/kernels/ops/layernorm/__init__.py

@@ -1,10 +1,11 @@
 """Layer-normalization kernels.
 
 Each operator is a :class:`~sglang.kernels.fused_op.BaseFusedOp` with a
-pure-``torch`` reference (``forward_native``) plus optimized CUDA backends,
+pure-``torch`` reference (``forward_native``) plus optimized per-device backends,
 all behind one signature. The public module-level functions are thin wrappers
-over module-level instances; auto-selection prefers the AOT ``sgl_kernel``
+over module-level instances; auto-selection follows the production default for
-implementation on CUDA and falls back to the native reference elsewhere.
+the live device: AOT ``sgl_kernel`` on CUDA, ``aiter`` (or rocm-triton for
+gemma) on ROCm, ``torch_npu`` on Ascend, native reference otherwise.
 Pick a specific backend with e.g.
 ``_RMSNORM.forward(x, w, backend=KernelBackend.JIT)`` or globally via
 ``SGLANG_FORCE_FUSED_OP_BACKEND``.
@@ -26,9 +27,21 @@ if TYPE_CHECKING:
 
 _NORM_DTYPES = ("float16", "bfloat16")
 _CUDA = frozenset({CapabilityRequirement.CUDA})
+_HIP = frozenset({CapabilityRequirement.HIP})
+_NPU = frozenset({CapabilityRequirement.NPU})
+# Unlike the gated-activation ops, sgl_kernel does *not* build the rmsnorm ops
+# for ROCm (production: ``if _is_cuda or _is_xpu or _is_musa: from sgl_kernel
+# import rmsnorm`` — HIP is absent), so AOT here is CUDA-only. ROCm instead has
+# an ``aiter`` path, and Ascend a ``torch_npu`` path — a clean illustration that
+# the same ``AOT`` provenance covers different devices per op.
+# Priority (best -> fallback) is device-agnostic; per-op CapabilityRequirement
+# decides eligibility, so on CUDA this resolves to AOT, on HIP to AITER, on NPU
+# to TORCH_NPU, each matching the production default for that device.
 _NORM_PRIORITY = (
     KernelBackend.AOT,
     KernelBackend.JIT,
+    KernelBackend.AITER,
+    KernelBackend.TORCH_NPU,
     KernelBackend.TORCH,
 )
 
@@ -44,6 +57,8 @@ class RMSNormOp(BaseFusedOp):
     capabilities = {
         KernelBackend.AOT: _CUDA,
         KernelBackend.JIT: _CUDA,
+        KernelBackend.AITER: _HIP,
+        KernelBackend.TORCH_NPU: _NPU,
     }
     format_signature = FormatSignature(
         supported_dtypes=_NORM_DTYPES,
@@ -52,6 +67,8 @@ class RMSNormOp(BaseFusedOp):
     descriptions = {
         KernelBackend.AOT: "RMS normalization (sgl_kernel wheel).",
         KernelBackend.JIT: "RMS normalization (sglang.jit_kernel).",
+        KernelBackend.AITER: "RMS normalization (aiter rmsnorm2d_fwd, ROCm).",
+        KernelBackend.TORCH_NPU: "RMS normalization (torch_npu, Ascend).",
         KernelBackend.TORCH: "RMS normalization (pure-torch reference).",
     }
 
@@ -103,6 +120,40 @@ class RMSNormOp(BaseFusedOp):
         jit_rmsnorm(input, weight, out, eps)
         return out
 
+    def forward_aiter(
+        self,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        out: Optional[torch.Tensor] = None,
+        enable_pdl: Optional[bool] = None,
+    ) -> torch.Tensor:
+        import torch
+        from aiter import rmsnorm2d_fwd
+
+        # Mirrors production srt/layers/layernorm.py: rmsnorm2d_fwd(out, x, w, eps)
+        # writes the normalized result in-place into ``out`` (ROCm path).
+        if out is None:
+            out = torch.empty_like(input)
+        rmsnorm2d_fwd(out, input, weight, eps)
+        return out
+
+    def forward_npu(
+        self,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        out: Optional[torch.Tensor] = None,
+        enable_pdl: Optional[bool] = None,
+    ) -> torch.Tensor:
+        import torch_npu
+
+        result = torch_npu.npu_rms_norm(input, weight, eps)[0]
+        if out is None:
+            return result
+        out.copy_(result)
+        return out
+
 
 class FusedAddRMSNormOp(BaseFusedOp):
     """In-place ``residual += input; input = RMSNorm(residual) * weight``.
@@ -116,6 +167,8 @@ class FusedAddRMSNormOp(BaseFusedOp):
     capabilities = {
         KernelBackend.AOT: _CUDA,
         KernelBackend.JIT: _CUDA,
+        KernelBackend.AITER: _HIP,
+        KernelBackend.TORCH_NPU: _NPU,
     }
     format_signature = FormatSignature(
         supported_dtypes=_NORM_DTYPES,
@@ -129,6 +182,10 @@ class FusedAddRMSNormOp(BaseFusedOp):
         KernelBackend.JIT: (
             "Fused residual-add + RMS normalization (sglang.jit_kernel)."
         ),
+        KernelBackend.AITER: ("Fused residual-add + RMS normalization (aiter, ROCm)."),
+        KernelBackend.TORCH_NPU: (
+            "Fused residual-add + RMS normalization (torch_npu, Ascend)."
+        ),
         KernelBackend.TORCH: (
             "Fused residual-add + RMS normalization (pure-torch reference)."
         ),
@@ -174,19 +231,65 @@ class FusedAddRMSNormOp(BaseFusedOp):
 
         return jit_fused_add_rmsnorm(input, residual, weight, eps)
 
+    def forward_aiter(
+        self,
+        input: torch.Tensor,
+        residual: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        enable_pdl: Optional[bool] = None,
+    ) -> None:
+        import torch
+        from aiter import rmsnorm2d_fwd_with_add
+
+        # aiter writes the normalized value and the new residual into separate
+        # out buffers (production call order: out, x, residual_out, residual, w,
+        # eps); copy them back to honor this op's in-place contract.
+        out = torch.empty_like(input)
+        residual_out = torch.empty_like(residual)
+        rmsnorm2d_fwd_with_add(out, input, residual_out, residual, weight, eps)
+        input.copy_(out)
+        residual.copy_(residual_out)
+
+    def forward_npu(
+        self,
+        input: torch.Tensor,
+        residual: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        enable_pdl: Optional[bool] = None,
+    ) -> None:
+        import torch_npu
+
+        # torch_npu.npu_add_rms_norm(residual, x, w, eps) -> (normed, _, new_sum)
+        out, _, residual_out = torch_npu.npu_add_rms_norm(residual, input, weight, eps)
+        input.copy_(out)
+        residual.copy_(residual_out)
+
 
 class GemmaRMSNormOp(BaseFusedOp):
     """``out = (input / RMS(input)) * (weight + 1)``; returns a tensor."""
 
     op = "layernorm.gemma_rmsnorm"
     priority = _NORM_PRIORITY
-    capabilities = {KernelBackend.AOT: _CUDA}
+    # AOT (sgl_kernel) on CUDA; JIT is the ROCm rocm-triton path
+    # (sglang.jit_kernel.minimax_m3) — a JIT provenance pinned to HIP, distinct
+    # from the CUDA-only JIT on the plain rmsnorm ops; torch_npu on Ascend.
+    capabilities = {
+        KernelBackend.AOT: _CUDA,
+        KernelBackend.JIT: _HIP,
+        KernelBackend.TORCH_NPU: _NPU,
+    }
     format_signature = FormatSignature(
         supported_dtypes=_NORM_DTYPES,
         description="out = (x / RMS(x)) * (weight + 1); returns tensor",
     )
     descriptions = {
         KernelBackend.AOT: "Gemma-style RMS normalization (sgl_kernel wheel).",
+        KernelBackend.JIT: (
+            "Gemma-style RMS normalization (rocm-triton, sglang.jit_kernel)."
+        ),
+        KernelBackend.TORCH_NPU: ("Gemma-style RMS normalization (torch_npu, Ascend)."),
         KernelBackend.TORCH: "Gemma-style RMS normalization (pure-torch reference).",
     }
 
@@ -221,13 +324,53 @@ class GemmaRMSNormOp(BaseFusedOp):
 
         return sgl_kernel.gemma_rmsnorm(input, weight, eps, out, enable_pdl)
 
+    def forward_jit(
+        self,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        out: Optional[torch.Tensor] = None,
+        enable_pdl: Optional[bool] = None,
+    ) -> torch.Tensor:
+        from sglang.jit_kernel.minimax_m3.rmsnorm import (
+            gemma_rmsnorm as rocm_triton_gemma_rmsnorm,
+        )
+
+        result = rocm_triton_gemma_rmsnorm(input, weight, eps)
+        if out is None:
+            return result
+        out.copy_(result)
+        return out
+
+    def forward_npu(
+        self,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        out: Optional[torch.Tensor] = None,
+        enable_pdl: Optional[bool] = None,
+    ) -> torch.Tensor:
+        import torch_npu
+
+        result = torch_npu.npu_gemma_rms_norm(input, weight, eps)[0]
+        if out is None:
+            return result
+        out.copy_(result)
+        return out
+
 
 class GemmaFusedAddRMSNormOp(BaseFusedOp):
     """In-place ``residual += input; input = GemmaRMSNorm(residual) * (weight + 1)``."""
 
     op = "layernorm.gemma_fused_add_rmsnorm"
     priority = _NORM_PRIORITY
-    capabilities = {KernelBackend.AOT: _CUDA}
+    # AOT (sgl_kernel) on CUDA; JIT is the ROCm rocm-triton path on HIP.
+    # NPU here would use ``sgl_kernel_npu.add_gemma_rms_norm`` (a distinct AOT-npu
+    # wheel provenance, not torch_npu) — deferred until that provenance lands.
+    capabilities = {
+        KernelBackend.AOT: _CUDA,
+        KernelBackend.JIT: _HIP,
+    }
     format_signature = FormatSignature(
         supported_dtypes=_NORM_DTYPES,
         in_place=True,
@@ -235,6 +378,10 @@ class GemmaFusedAddRMSNormOp(BaseFusedOp):
     )
     descriptions = {
         KernelBackend.AOT: ("Gemma-style fused residual-add + RMS normalization."),
+        KernelBackend.JIT: (
+            "Gemma-style fused residual-add + RMS normalization "
+            "(rocm-triton, sglang.jit_kernel)."
+        ),
         KernelBackend.TORCH: (
             "Gemma-style fused residual-add + RMS normalization "
             "(pure-torch reference)."
@@ -271,6 +418,25 @@ class GemmaFusedAddRMSNormOp(BaseFusedOp):
             input, residual, weight, eps, enable_pdl
         )
 
+    def forward_jit(
+        self,
+        input: torch.Tensor,
+        residual: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        enable_pdl: Optional[bool] = None,
+    ) -> None:
+        from sglang.jit_kernel.minimax_m3.rmsnorm import (
+            gemma_fused_add_rmsnorm as rocm_triton_gemma_fused_add_rmsnorm,
+        )
+
+        # rocm-triton returns (normed, new_residual); honor the in-place contract.
+        norm_out, residual_out = rocm_triton_gemma_fused_add_rmsnorm(
+            input, residual, weight, eps
+        )
+        input.copy_(norm_out)
+        residual.copy_(residual_out)
+
 
 _RMSNORM = register_fused_op(RMSNormOp(), __name__, "_RMSNORM")
 _FUSED_ADD_RMSNORM = register_fused_op(

python/sglang/kernels/spec.py

@@ -33,7 +33,8 @@ class KernelBackend(str, Enum):
     (the ``sgl_kernel`` wheel, built for CUDA *and* ROCm) are both cross-device;
     which devices a given op supports is expressed by its
     :class:`CapabilityRequirement` list. Platform-specific libraries (e.g.
-    ``aiter`` on AMD) are just additional provenance values.
+    ``aiter`` on AMD, ``torch_npu`` on Ascend) are just additional provenance
+    values, each pinned to its device by its ``CapabilityRequirement``.
     """
 
     TORCH = "torch"  # pure-torch reference (forward_native)
@@ -45,7 +46,8 @@ class KernelBackend(str, Enum):
     FLASHINFER = "flashinfer"
     DEEPGEMM = "deepgemm"
     AITER = "aiter"  # AMD aiter library (device=HIP)
-    # TODO(RFC #29630): more provenance as needed (npu / cpu-avx, ...)
+    TORCH_NPU = "torch_npu"  # Ascend NPU vendor runtime (device=NPU)
+    # TODO(RFC #29630): more provenance as needed (cpu-avx, sgl_kernel_npu, ...)
 
 
 class DeviceType(str, Enum):
@@ -53,8 +55,9 @@ class DeviceType(str, Enum):
 
     CUDA = "cuda"
     HIP = "hip"
+    NPU = "npu"  # Ascend NPU (torch_npu / sgl_kernel_npu)
     CPU = "cpu"
-    # TODO(RFC #29630): NPU / XPU / ... as backends land.
+    # TODO(RFC #29630): XPU / MUSA / ... as backends land.
 
 
 class PlatformInfo(msgspec.Struct, frozen=True):
@@ -98,6 +101,9 @@ class PlatformInfo(msgspec.Struct, frozen=True):
         try:
             if torch.version.hip is not None and torch.cuda.is_available():
                 return cls(device_type="hip")
+            npu = getattr(torch, "npu", None)
+            if npu is not None and npu.is_available():
+                return cls(device_type="npu")
             if torch.cuda.is_available():
                 major, minor = torch.cuda.get_device_capability()
                 return cls(
@@ -123,7 +129,7 @@ class CapabilityRequirement(msgspec.Struct, frozen=True):
     (``min_cuda_arch`` / ``max_cuda_arch`` apply only when ``device == CUDA``).
 
     The device-only cases are so common that they are exposed as class constants
-    (``CapabilityRequirement.CUDA`` / ``.HIP``); use :meth:`cuda` for an
+    (``CapabilityRequirement.CUDA`` / ``.HIP`` / ``.NPU``); use :meth:`cuda` for an
     arch-bounded CUDA requirement (e.g. ``CapabilityRequirement.cuda(
     min_sm=(10, 0))`` for SM100+).
     """
@@ -136,6 +142,7 @@ class CapabilityRequirement(msgspec.Struct, frozen=True):
     # instances of the class itself). ClassVar keeps them out of msgspec fields.
     CUDA: ClassVar[CapabilityRequirement]
     HIP: ClassVar[CapabilityRequirement]
+    NPU: ClassVar[CapabilityRequirement]
 
     @classmethod
     def cuda(
@@ -160,6 +167,7 @@ class CapabilityRequirement(msgspec.Struct, frozen=True):
 
 CapabilityRequirement.CUDA = CapabilityRequirement(device=DeviceType.CUDA)
 CapabilityRequirement.HIP = CapabilityRequirement(device=DeviceType.HIP)
+CapabilityRequirement.NPU = CapabilityRequirement(device=DeviceType.NPU)
 
 
 def capabilities_satisfied(

test/registered/kernels/test_fused_op.py

@@ -152,6 +152,8 @@ class TestBaseFusedOp(unittest.TestCase):
                 KernelBackend.TORCH_COMPILE,
                 KernelBackend.JIT,
                 KernelBackend.AOT,
+                KernelBackend.AITER,
+                KernelBackend.TORCH_NPU,
             },
         )
         # Dotted targets resolve to the bound backend methods.

test/registered/kernels/test_kernels_namespace.py

@@ -28,15 +28,37 @@ EXPECTED_OPS = {
         "torch",
         "torch_compile",
     },
-    "layernorm.rmsnorm": {"aot", "jit", "torch", "torch_compile"},
+    "activation.relu2": {"jit", "torch", "torch_compile"},
+    "activation.gelu_quick": {"aot", "torch", "torch_compile"},
+    "layernorm.rmsnorm": {
+        "aot",
+        "jit",
+        "aiter",
+        "torch_npu",
+        "torch",
+        "torch_compile",
+    },
     "layernorm.fused_add_rmsnorm": {
+        "aot",
+        "jit",
+        "aiter",
+        "torch_npu",
+        "torch",
+        "torch_compile",
+    },
+    "layernorm.gemma_rmsnorm": {
+        "aot",
+        "jit",
+        "torch_npu",
+        "torch",
+        "torch_compile",
+    },
+    "layernorm.gemma_fused_add_rmsnorm": {
         "aot",
         "jit",
         "torch",
         "torch_compile",
     },
-    "layernorm.gemma_rmsnorm": {"aot", "torch", "torch_compile"},
-    "layernorm.gemma_fused_add_rmsnorm": {"aot", "torch", "torch_compile"},
     # curated dual/single-backend wrapper ops
     "gemm.fp8_scaled_mm": {"aot"},
     "gemm.dsv3_fused_a_gemm": {"aot", "jit"},
@@ -249,6 +271,41 @@ class TestKernelsNamespace(unittest.TestCase):
         finally:
             fo._platform = saved
 
+    def test_layernorm_cross_device_coverage(self):
+        # The rmsnorm ops illustrate that the *same* provenance covers different
+        # devices per op: AOT (sgl_kernel) is CUDA-only here (sgl_kernel does not
+        # build rmsnorm for ROCm), so HIP falls to AITER and NPU to torch_npu,
+        # each matching the production default for that device. gemma uses a
+        # rocm-triton JIT path on HIP -- a JIT provenance pinned to HIP, unlike
+        # the CUDA-only JIT on plain rmsnorm.
+        import sglang.kernels.fused_op as fo
+        from sglang.kernels.ops.layernorm import (
+            _FUSED_ADD_RMSNORM,
+            _GEMMA_RMSNORM,
+            _RMSNORM,
+        )
+
+        B = self.K.KernelBackend
+        cuda = self.K.PlatformInfo(device_type="cuda", cuda_arch_major=9)
+        hip = self.K.PlatformInfo(device_type="hip")
+        npu = self.K.PlatformInfo(device_type="npu")
+        saved = fo._platform
+        try:
+            for plat, expect in ((cuda, B.AOT), (hip, B.AITER), (npu, B.TORCH_NPU)):
+                fo._platform = lambda p=plat: p
+                self.assertEqual(_RMSNORM._resolve_backend(), expect)
+                self.assertEqual(_FUSED_ADD_RMSNORM._resolve_backend(), expect)
+            # gemma: AOT on CUDA, rocm-triton JIT on HIP, torch_npu on NPU.
+            for plat, expect in ((cuda, B.AOT), (hip, B.JIT), (npu, B.TORCH_NPU)):
+                fo._platform = lambda p=plat: p
+                self.assertEqual(_GEMMA_RMSNORM._resolve_backend(), expect)
+            # AOT rmsnorm is CUDA-only (not HIP) -- distinct from activation's AOT.
+            fo._platform = lambda: hip
+            self.assertFalse(_RMSNORM.backend_eligible(B.AOT))
+            self.assertTrue(_RMSNORM.backend_eligible(B.AITER))
+        finally:
+            fo._platform = saved
+
     def test_selector_explicit_backend(self):
         spec = self.K.select_kernel(
             "layernorm.rmsnorm", backend=self.K.KernelBackend.JIT
@@ -301,6 +358,7 @@ class TestKernelsNamespace(unittest.TestCase):
         # and dedup, so {CUDA, HIP} == {HIP, CUDA}.
         self.assertEqual(cap.CUDA, cap(device=dev.CUDA))
         self.assertEqual(cap.HIP, cap(device=dev.HIP))
+        self.assertEqual(cap.NPU, cap(device=dev.NPU))
         self.assertEqual({cap.CUDA, cap.HIP}, {cap.HIP, cap.CUDA})
         self.assertEqual(len({cap.CUDA, cap(device=dev.CUDA)}), 1)
         # cuda(min_sm=...) factory: an SM100+ CUDA requirement.