[Diffusion] Enable shared RMSNorm dispatch for SenseNova-U1 (#39705)

Signed-off-by: syd520zy <529477025@qq.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>

python/sglang/multimodal_gen/runtime/models/sensenova_u1/neo_unify/modeling_qwen3.py

@@ -29,7 +29,6 @@ from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, can_return_tuple
 from transformers.utils.deprecation import deprecate_kwarg
 
-from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.srt.layers.layernorm import RMSNorm
 
 from .transformers_compat import (
@@ -347,7 +346,6 @@ def make_qwen3_rms_norm(hidden_size: int, eps: float) -> RMSNorm:
         hidden_size,
         eps=eps,
         cast_x_before_out_mul=True,
-        force_native=not current_platform.is_npu(),
     )
 
 

python/sglang/multimodal_gen/test/unit/test_sensenova_u1.py

@@ -380,18 +380,12 @@ def test_sensenova_u1_npu_fia_checks_operator_availability(monkeypatch, availabl
     assert npu_fia_available() is available
 
 
-@pytest.mark.parametrize(
-    ("is_npu", "uses_native"),
-    [(False, True), (True, False)],
-)
-def test_sensenova_u1_shared_rmsnorm_dispatch(monkeypatch, is_npu, uses_native):
-    monkeypatch.setattr(current_platform, "is_npu", lambda: is_npu)
-
+def test_sensenova_u1_shared_rmsnorm_uses_framework_dispatch():
     norm = make_qwen3_rms_norm(64, eps=1e-6)
 
     assert isinstance(norm, RMSNorm)
     assert norm.cast_x_before_out_mul
-    assert (norm._forward_method == norm.forward_native) is uses_native
+    assert norm._forward_method != norm.forward_native
 
 
 @torch.no_grad()

python/sglang/srt/layers/layernorm.py

@@ -637,6 +637,11 @@ class RMSNorm(BaseFusedOp):
             if residual is not None:
                 return x, residual
             return x
+        if not x.is_cuda:
+            # AITER kernels dereference activations as GPU pointers; a CPU
+            # input (e.g. unit tests building modules on CPU) aborts the
+            # process with HSA_STATUS_ERROR_MEMORY_FAULT instead of raising.
+            return self.forward_native(x, residual, post_residual_addition)
         if self.weight.data.dtype != x.dtype:
             # AITER's ROCm rmsnorm2d_fwd requires weight/activation dtypes to match;
             # FP32 weight + BF16 activation yields finite-but-corrupted output on gfx950.