diff --git a/docs/docs/hardware-platforms/ascend-npus/getting-started/installation.mdx b/docs/docs/hardware-platforms/ascend-npus/getting-started/installation.mdx index 2de76d2973aa..d805b4328eda 100644 --- a/docs/docs/hardware-platforms/ascend-npus/getting-started/installation.mdx +++ b/docs/docs/hardware-platforms/ascend-npus/getting-started/installation.mdx @@ -252,6 +252,18 @@ For installation of Triton on Ascend nightly builds or from sources, follow [ins We provide SGL kernels for Ascend NPU, check [installation guide](https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/sgl_kernel_npu/README.md). + +Gemma-family models (including Qwen3.5) call `sgl_kernel_npu.norm.gemma_rmsnorm`, +whose implementation is chosen when the wheel is built: the native `torch_npu` +operator for Ascend 910B/910C, and ACLNN RMSNorm for Ascend 950, which does not +register that operator. **Ascend 950 users must install the `950` release +package** (`950` appears in the asset name), or build from source with +`bash build.sh -a kernels Ascend950PR_9599`, the target that package is built +with. On Ascend 950, a `910b` or `a3` package fails on the first Gemma forward, +and so does a package from a release that predates this module; check with +`python -c "import sgl_kernel_npu.norm.gemma_rmsnorm"`. + + #### DeepEP-compatible Library We provide a DeepEP-compatible Library as a drop-in replacement of deepseek-ai's DeepEP library, check the [installation guide](https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/deep_ep/README.md). diff --git a/python/sglang/kernels/fused_op.py b/python/sglang/kernels/fused_op.py index fb1d15bb5c4b..d7c5b33d1153 100644 --- a/python/sglang/kernels/fused_op.py +++ b/python/sglang/kernels/fused_op.py @@ -112,6 +112,7 @@ KernelBackend.FLASHINFER: "forward_flashinfer", KernelBackend.DEEPGEMM: "forward_deepgemm", KernelBackend.AITER: "forward_aiter", + KernelBackend.SGL_KERNEL_NPU: "forward_sgl_kernel_npu", KernelBackend.TORCH_NPU: "forward_torch_npu", } @@ -130,6 +131,7 @@ KernelBackend.DEEPGEMM, KernelBackend.CUTE_DSL, KernelBackend.AITER, + KernelBackend.SGL_KERNEL_NPU, KernelBackend.TORCH_NPU, KernelBackend.TRITON, KernelBackend.TORCH, @@ -448,6 +450,9 @@ def forward_deepgemm(self, *args, **kwargs): def forward_aiter(self, *args, **kwargs): raise NotImplementedError(f"{self._op_label()}: no aiter backend") + def forward_sgl_kernel_npu(self, *args, **kwargs): + raise NotImplementedError(f"{self._op_label()}: no sgl_kernel_npu backend") + def forward_torch_npu(self, *args, **kwargs): raise NotImplementedError(f"{self._op_label()}: no torch_npu backend") diff --git a/python/sglang/kernels/ops/layernorm/__init__.py b/python/sglang/kernels/ops/layernorm/__init__.py index 1e71b01c92ef..e0ec533fd562 100644 --- a/python/sglang/kernels/ops/layernorm/__init__.py +++ b/python/sglang/kernels/ops/layernorm/__init__.py @@ -5,7 +5,8 @@ all behind one signature. The public module-level functions are thin wrappers over module-level instances; auto-selection follows the production default for the live device: AOT ``sgl_kernel`` on CUDA, ``aiter`` (or rocm-triton for -gemma) on ROCm, ``torch_npu`` on Ascend, native reference otherwise. +gemma) on ROCm, ``torch_npu`` on Ascend (``sgl_kernel_npu`` for gemma, whose +implementation is SoC-specific), 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``. @@ -13,6 +14,7 @@ from __future__ import annotations +import importlib from typing import TYPE_CHECKING, Optional from sglang.kernels.fused_op import BaseFusedOp, register_fused_op @@ -35,17 +37,36 @@ # 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. +# decides eligibility, so on CUDA this resolves to AOT, on HIP to AITER, and on +# NPU to the provider implemented by each operator. _NORM_PRIORITY = ( KernelBackend.AOT, KernelBackend.JIT, KernelBackend.AITER, + KernelBackend.SGL_KERNEL_NPU, KernelBackend.TORCH_NPU, KernelBackend.TORCH, ) +def _sgl_kernel_npu_gemma(module: str, symbol: str): + """Resolve a Gemma kernel from ``sgl_kernel_npu``, or explain what is missing. + + The Gemma provider is picked when the sgl-kernel-npu wheel is built (native + ``torch_npu`` operator on Ascend 910, ACLNN on Ascend 950), so an older or + mismatched wheel shows up here as a plain ImportError. Unlike + ``srt/layers/layernorm.py``, this path does not fall back to ``torch_npu``: + a backend selected by name must not silently run a different provenance. + """ + try: + return getattr(importlib.import_module(module), symbol) + except (ImportError, AttributeError) as error: + raise RuntimeError( + "Gemma RMSNorm on Ascend requires a target-specific sgl-kernel-npu " + f"wheel that provides {module}.{symbol}" + ) from error + + class RMSNormOp(BaseFusedOp): """``out = (input / RMS(input)) * weight``; returns a tensor. @@ -276,11 +297,11 @@ class GemmaRMSNormOp(BaseFusedOp): priority = _NORM_PRIORITY # AOT (sgl_kernel) on CUDA; JIT is the ROCm rocm-triton path # (sglang.kernels.ops.moe.minimax_m3_swiglu) — a JIT provenance pinned to HIP, distinct - # from the CUDA-only JIT on the plain rmsnorm ops; torch_npu on Ascend. + # from the CUDA-only JIT on the plain rmsnorm ops; sgl_kernel_npu on Ascend. capabilities = { KernelBackend.AOT: _CUDA, KernelBackend.JIT: _HIP, - KernelBackend.TORCH_NPU: _NPU, + KernelBackend.SGL_KERNEL_NPU: _NPU, } format_signature = FormatSignature( supported_dtypes=_NORM_DTYPES, @@ -291,7 +312,9 @@ class GemmaRMSNormOp(BaseFusedOp): KernelBackend.JIT: ( "Gemma-style RMS normalization (rocm-triton, sglang.kernels.jit)." ), - KernelBackend.TORCH_NPU: ("Gemma-style RMS normalization (torch_npu, Ascend)."), + KernelBackend.SGL_KERNEL_NPU: ( + "Gemma-style RMS normalization (sgl_kernel_npu, Ascend)." + ), KernelBackend.TORCH: "Gemma-style RMS normalization (pure-torch reference).", } @@ -344,7 +367,7 @@ def forward_jit( out.copy_(result) return out - def forward_torch_npu( + def forward_sgl_kernel_npu( self, input: torch.Tensor, weight: torch.Tensor, @@ -352,9 +375,10 @@ def forward_torch_npu( 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] + npu_gemma_rms_norm = _sgl_kernel_npu_gemma( + "sgl_kernel_npu.norm.gemma_rmsnorm", "npu_gemma_rms_norm" + ) + result, _ = npu_gemma_rms_norm(input, weight, eps) if out is None: return result out.copy_(result) @@ -367,11 +391,12 @@ class GemmaFusedAddRMSNormOp(BaseFusedOp): op = "layernorm.gemma_fused_add_rmsnorm" priority = _NORM_PRIORITY # 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. + # NPU uses ``sgl_kernel_npu.norm.add_rmsnorm_bias.add_gemma_rms_norm`` — the + # SGL_KERNEL_NPU provenance, not torch_npu. capabilities = { KernelBackend.AOT: _CUDA, KernelBackend.JIT: _HIP, + KernelBackend.SGL_KERNEL_NPU: _NPU, } format_signature = FormatSignature( supported_dtypes=_NORM_DTYPES, @@ -384,6 +409,10 @@ class GemmaFusedAddRMSNormOp(BaseFusedOp): "Gemma-style fused residual-add + RMS normalization " "(rocm-triton, sglang.kernels.jit)." ), + KernelBackend.SGL_KERNEL_NPU: ( + "Gemma-style fused residual-add + RMS normalization " + "(sgl_kernel_npu, Ascend)." + ), KernelBackend.TORCH: ( "Gemma-style fused residual-add + RMS normalization (pure-torch reference)." ), @@ -438,6 +467,22 @@ def forward_jit( input.copy_(norm_out) residual.copy_(residual_out) + def forward_sgl_kernel_npu( + self, + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + enable_pdl: Optional[bool] = None, + ) -> None: + add_gemma_rms_norm = _sgl_kernel_npu_gemma( + "sgl_kernel_npu.norm.add_rmsnorm_bias", "add_gemma_rms_norm" + ) + # sgl_kernel_npu returns (normed, new_residual); honor the in-place contract. + norm_out, residual_out = add_gemma_rms_norm(input, weight, residual, eps) + input.copy_(norm_out) + residual.copy_(residual_out) + _RMSNORM = register_fused_op(RMSNormOp(), __name__, "_RMSNORM") _FUSED_ADD_RMSNORM = register_fused_op( diff --git a/python/sglang/kernels/spec.py b/python/sglang/kernels/spec.py index 913534a26525..05eaa0da0313 100644 --- a/python/sglang/kernels/spec.py +++ b/python/sglang/kernels/spec.py @@ -48,8 +48,9 @@ class KernelBackend(str, Enum): FLASHINFER = "flashinfer" DEEPGEMM = "deepgemm" AITER = "aiter" # AMD aiter library (device=HIP) + SGL_KERNEL_NPU = "sgl_kernel_npu" # SGLang Ascend kernel package (device=NPU) TORCH_NPU = "torch_npu" # Ascend NPU vendor runtime (device=NPU) - # TODO(RFC #29630): more provenance as needed (cpu-avx, sgl_kernel_npu, ...) + # TODO(RFC #29630): more provenance as needed (cpu-avx, ...) class DeviceType(str, Enum): diff --git a/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py index 3c905bf80da1..b2dab4d53a2e 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, List, Optional, Tuple import torch +import torch.nn.functional as F from torch.nn.parameter import Parameter from sglang.srt.hardware_backend.npu.utils import NPUACLFormat, npu_format_cast @@ -204,10 +205,15 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: weight = layer.weight.data if weight.dtype == torch.float8_e4m3fn: # Offline (ModelSlim) path: weight is already MXFP8-quantised and - # layer.weight_scale holds the uint8 block scales [out, in/32]. Only - # re-layout to [in, out] / [in//64, out, 2] strided views below. + # layer.weight_scale holds one uint8 scale per 32 input elements. + # Pair the scales for the NPU kernel, padding an odd final count + # (e.g. K=4304 -> 135 scales -> 136 -> 68 pairs). n_dim, k_dim = layer.weight_scale.data.shape - scale = layer.weight_scale.data.reshape(n_dim, k_dim // 2, 2) + scale_data = layer.weight_scale.data + if k_dim % 2 != 0: + scale_data = F.pad(scale_data, (0, 1), mode="constant", value=0) + k_dim += 1 + scale = scale_data.reshape(n_dim, k_dim // 2, 2) layer.weight = Parameter(weight.transpose(0, 1), requires_grad=False) layer.weight_scale_inv = Parameter( scale.transpose(0, 1), requires_grad=False diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 6068e8f2095f..c88791e07d35 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -184,6 +184,19 @@ def is_supported_rmsnorm_hf_hidden_size(d: int) -> bool: import torch_npu from sgl_kernel_npu.norm.add_rmsnorm_bias import add_gemma_rms_norm + try: + from sgl_kernel_npu.norm.gemma_rmsnorm import npu_gemma_rms_norm + except ImportError: + # sgl-kernel-npu wheels built before the target-specific Gemma provider + # landed expose only the torch_npu operator, which is exactly what that + # provider binds to on A2/A3 — so those deployments keep working + # unchanged. On Ascend 950 the operator is unregistered and the first Gemma + # forward fails loudly, same as before this indirection existed. The + # kernels registry (sglang/kernels/ops/layernorm) deliberately does not + # fall back: a backend selected by name there must not silently run a + # different provenance. + npu_gemma_rms_norm = torch_npu.npu_gemma_rms_norm + _NPU_GEMMA_RMS_NORM_TRITON_MAX_HIDDEN_SIZE = 5120 @@ -1248,7 +1261,7 @@ def forward_npu( ) return norm_out, residual - x, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.variance_epsilon) + x, _ = npu_gemma_rms_norm(x, self.weight, self.variance_epsilon) return x def forward_xpu( @@ -1370,7 +1383,7 @@ def forward_hip(self, x, residual: Optional[torch.Tensor] = None): def forward_npu(self, x, residual: Optional[torch.Tensor] = None): if residual is not None: return self.forward_native(x, residual) - output, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.eps) + output, _ = npu_gemma_rms_norm(x, self.weight, self.eps) return output def extra_repr(self): diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp8.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp8.py index 8c82ed93293f..d82f977bc398 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp8.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp8.py @@ -60,11 +60,15 @@ def create_weights( ) layer.register_parameter("weight", weight) - # msmodelslim exports weight_scale as uint8, shape [out, in/32]. + # msmodelslim exports one uint8 scale per 32 input elements. Keep the + # final partial block: Qwen3-VL has visual projections such as K=4304, + # whose checkpoint scale shape is [out, ceil(4304/32)] = [out, 135]. # NOTE: Named "weight_scale" (not "weight_scale_inv") to match the # checkpoint key exported by msmodelslim; the kernel re-layouts it into # weight_scale_inv during process_weights_after_loading. - scale_dim = input_size_per_partition // MXFP8_BLOCK_SIZE + scale_dim = ( + input_size_per_partition + MXFP8_BLOCK_SIZE - 1 + ) // MXFP8_BLOCK_SIZE weight_scale = GroupQuantScaleParameter( data=torch.empty( (output_size_per_partition, scale_dim), diff --git a/test/registered/unit/kernels/test_kernels_namespace.py b/test/registered/unit/kernels/test_kernels_namespace.py index 41f47e9a1963..aa1bae6081aa 100644 --- a/test/registered/unit/kernels/test_kernels_namespace.py +++ b/test/registered/unit/kernels/test_kernels_namespace.py @@ -91,7 +91,7 @@ def test_activation_default_backend(monkeypatch, device, expect): ("_RMSNORM", "npu", "torch_npu"), ("_GEMMA_RMSNORM", "cuda", "aot"), ("_GEMMA_RMSNORM", "hip", "jit"), # rocm-triton JIT pinned to HIP - ("_GEMMA_RMSNORM", "npu", "torch_npu"), + ("_GEMMA_RMSNORM", "npu", "sgl_kernel_npu"), ], ) def test_layernorm_default_backend(monkeypatch, op_attr, device, expect): diff --git a/test/registered/unit/layers/quantization/test_modelslim_mxfp8.py b/test/registered/unit/layers/quantization/test_modelslim_mxfp8.py new file mode 100644 index 000000000000..60e4ed1c734a --- /dev/null +++ b/test/registered/unit/layers/quantization/test_modelslim_mxfp8.py @@ -0,0 +1,61 @@ +"""CPU regression tests for ModelSlim MXFP8 weight-scale layouts.""" + +import unittest + +import torch + +from sglang.srt.layers.quantization.modelslim.schemes.modelslim_mxfp8 import ( + ModelSlimMXFP8Scheme, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class TestModelSlimMXFP8ScaleLayout(CustomTestCase): + def setUp(self): + super().setUp() + self.scheme = ModelSlimMXFP8Scheme() + + def test_scale_placeholder_rounds_up_partial_block(self): + layer = torch.nn.Module() + + self.scheme.create_weights( + layer=layer, + input_size_per_partition=4304, + output_partition_sizes=[2], + input_size=4304, + output_size=2, + params_dtype=torch.bfloat16, + ) + + self.assertEqual(layer.weight_scale.shape, (2, 135)) + + def test_post_load_pads_odd_scale_count_for_pair_layout(self): + layer = torch.nn.Module() + layer.register_parameter( + "weight", + torch.nn.Parameter( + torch.empty((2, 4304), dtype=torch.float8_e4m3fn), + requires_grad=False, + ), + ) + layer.register_parameter( + "weight_scale", + torch.nn.Parameter( + torch.ones((2, 135), dtype=torch.uint8), requires_grad=False + ), + ) + layer.register_parameter("bias", None) + + self.scheme.process_weights_after_loading(layer) + + self.assertEqual(layer.weight.shape, (4304, 2)) + self.assertEqual(layer.weight_scale_inv.shape, (68, 2, 2)) + self.assertTrue(torch.all(layer.weight_scale_inv[-1, :, 1] == 0)) + self.assertFalse(hasattr(layer, "weight_scale")) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/npu/test_npu_gemma_rmsnorm.py b/test/registered/unit/npu/test_npu_gemma_rmsnorm.py new file mode 100644 index 000000000000..ae15b28a14e7 --- /dev/null +++ b/test/registered/unit/npu/test_npu_gemma_rmsnorm.py @@ -0,0 +1,187 @@ +import re +import sys +from pathlib import Path +from types import ModuleType +from unittest.mock import MagicMock, patch + +import pytest +import torch + +import sglang.kernels.fused_op as fused_op_module +from sglang import kernels +from sglang.kernels.ops.layernorm import ( + GemmaFusedAddRMSNormOp, + GemmaRMSNormOp, +) +from sglang.kernels.spec import CapabilityRequirement, KernelBackend, PlatformInfo +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +def _fake_sgl_kernel_npu(gemma_kernel, add_gemma_kernel): + package = ModuleType("sgl_kernel_npu") + package.__path__ = [] + norm_package = ModuleType("sgl_kernel_npu.norm") + norm_package.__path__ = [] + gemma_module = ModuleType("sgl_kernel_npu.norm.gemma_rmsnorm") + gemma_module.npu_gemma_rms_norm = gemma_kernel + add_module = ModuleType("sgl_kernel_npu.norm.add_rmsnorm_bias") + add_module.add_gemma_rms_norm = add_gemma_kernel + return { + "sgl_kernel_npu": package, + "sgl_kernel_npu.norm": norm_package, + "sgl_kernel_npu.norm.gemma_rmsnorm": gemma_module, + "sgl_kernel_npu.norm.add_rmsnorm_bias": add_module, + } + + +@pytest.mark.parametrize( + "op_name", + ["layernorm.gemma_rmsnorm", "layernorm.gemma_fused_add_rmsnorm"], +) +def test_registry_exposes_only_sgl_kernel_npu_provider(op_name): + registered_backends = {spec.backend for spec in kernels.registry.get(op_name)} + spec = kernels.registry.get_backend(op_name, KernelBackend.SGL_KERNEL_NPU) + + assert KernelBackend.TORCH_NPU not in registered_backends + assert spec.capabilities == frozenset({CapabilityRequirement.NPU}) + assert spec.target.endswith(".forward_sgl_kernel_npu") + + +@pytest.mark.parametrize("op_cls", [GemmaRMSNormOp, GemmaFusedAddRMSNormOp]) +def test_sgl_kernel_npu_is_selected_on_npu(op_cls): + with patch.object( + fused_op_module, "_platform", return_value=PlatformInfo(device_type="npu") + ): + backend = op_cls().auto_selected_backend() + + assert backend is KernelBackend.SGL_KERNEL_NPU + + +def test_missing_kernel_package_reports_actionable_error(): + package = ModuleType("sgl_kernel_npu") + package.__path__ = [] + norm_package = ModuleType("sgl_kernel_npu.norm") + norm_package.__path__ = [] + + with patch.dict( + sys.modules, + { + "sgl_kernel_npu": package, + "sgl_kernel_npu.norm": norm_package, + }, + ): + sys.modules.pop("sgl_kernel_npu.norm.gemma_rmsnorm", None) + with pytest.raises(RuntimeError, match="requires a target-specific"): + GemmaRMSNormOp().forward_sgl_kernel_npu( + torch.randn(2, 4), torch.randn(4), 1e-5 + ) + + +def test_incompatible_kernel_package_reports_actionable_error(): + """A wheel that imports but lacks the symbol must name the module it lacks. + + Covers the AttributeError branch (the test above covers ImportError), and + pins each op to its own module so the message cannot drift into a + copy-pasted sibling name. + """ + modules = _fake_sgl_kernel_npu(MagicMock(), MagicMock()) + del modules["sgl_kernel_npu.norm.gemma_rmsnorm"].npu_gemma_rms_norm + del modules["sgl_kernel_npu.norm.add_rmsnorm_bias"].add_gemma_rms_norm + + with patch.dict(sys.modules, modules): + with pytest.raises(RuntimeError, match=r"norm\.gemma_rmsnorm\."): + GemmaRMSNormOp().forward_sgl_kernel_npu( + torch.randn(2, 4), torch.randn(4), 1e-5 + ) + with pytest.raises(RuntimeError, match=r"norm\.add_rmsnorm_bias\."): + GemmaFusedAddRMSNormOp().forward_sgl_kernel_npu( + torch.randn(2, 4), torch.randn(2, 4), torch.randn(4), 1e-5 + ) + + +def test_sgl_kernel_npu_normal_out_contract(): + x = torch.randn(2, 4) + weight = torch.randn(4) + expected = torch.randn_like(x) + out = torch.empty_like(x) + gemma_kernel = MagicMock(return_value=(expected, "rstd")) + + with patch.dict(sys.modules, _fake_sgl_kernel_npu(gemma_kernel, MagicMock())): + result = GemmaRMSNormOp().forward_sgl_kernel_npu(x, weight, 1e-5, out=out) + + assert result is out + torch.testing.assert_close(out, expected) + gemma_kernel.assert_called_once_with(x, weight, 1e-5) + + +def test_sgl_kernel_npu_fused_in_place_contract(): + x = torch.randn(2, 4) + residual = torch.randn(2, 4) + weight = torch.randn(4) + norm_output = torch.randn_like(x) + residual_sum = torch.randn_like(residual) + add_gemma_kernel = MagicMock(return_value=(norm_output, residual_sum)) + + with patch.dict(sys.modules, _fake_sgl_kernel_npu(MagicMock(), add_gemma_kernel)): + result = GemmaFusedAddRMSNormOp().forward_sgl_kernel_npu( + x, residual, weight, 1e-5 + ) + + assert result is None + torch.testing.assert_close(x, norm_output) + torch.testing.assert_close(residual, residual_sum) + add_gemma_kernel.assert_called_once() + args = add_gemma_kernel.call_args.args + assert args[0] is x + assert args[1] is weight + assert args[2] is residual + assert args[3] == 1e-5 + + +@pytest.mark.parametrize( + ("layer_name", "eps_attr"), + [("GemmaRMSNorm", "variance_epsilon"), ("Gemma3RMSNorm", "eps")], +) +def test_srt_gemma_layers_delegate_plain_npu_path(layer_name, eps_attr): + from sglang.srt.layers import layernorm as layernorm_module + + layer = getattr(layernorm_module, layer_name)(4) + x = torch.randn(2, 4) + unified_op = MagicMock(return_value=(x, "rstd")) + + with patch.object(layernorm_module, "npu_gemma_rms_norm", unified_op, create=True): + result = layer.forward_npu(x) + + assert result is x + unified_op.assert_called_once_with(x, layer.weight, getattr(layer, eps_attr)) + + +def test_srt_falls_back_to_torch_npu_on_wheels_without_the_provider(): + """The srt-layer provider import must stay guarded, with a torch_npu fallback. + + ``srt/layers/layernorm.py`` imports the provider at module scope under + ``if _is_npu``, which CPU CI never executes -- hence a source-shape check. + Turning it back into a hard import would break every NPU deployment running + an sgl-kernel-npu wheel from before the staged Gemma provider, at + ``import sglang.srt.layers.layernorm`` time and for all models, not just + Gemma ones. + """ + from sglang.srt.layers import layernorm as layernorm_module + + source = Path(layernorm_module.__file__).read_text(encoding="utf-8") + guarded_import = re.search( + r"try:\s*\n" + r"\s*from sgl_kernel_npu\.norm\.gemma_rmsnorm import npu_gemma_rms_norm\s*\n" + r"\s*except ImportError:\s*\n" + r"(?:\s*#.*\n)*" + r"\s*npu_gemma_rms_norm = torch_npu\.npu_gemma_rms_norm\s*\n", + source, + ) + + assert guarded_import is not None + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))