From 791c7850d0960fd768102f71e7d999b036bb75ba Mon Sep 17 00:00:00 2001
From: faceless void <100186390+syd520zy@users.noreply.github.com>
Date: Sun, 20 Sep 2026 22:08:02 +0800
Subject: [PATCH] [Diffusion] Enable shared RMSNorm dispatch for SenseNova-U1
 (#39705)

Signed-off-by: syd520zy <529477025@qq.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>
---
 .../models/sensenova_u1/neo_unify/modeling_qwen3.py    |  2 --
 .../multimodal_gen/test/unit/test_sensenova_u1.py      | 10 ++--------
 python/sglang/srt/layers/layernorm.py                  |  5 +++++
 3 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/models/sensenova_u1/neo_unify/modeling_qwen3.py b/python/sglang/multimodal_gen/runtime/models/sensenova_u1/neo_unify/modeling_qwen3.py
index 68d4d5396..bb5c79a82 100644
--- a/python/sglang/multimodal_gen/runtime/models/sensenova_u1/neo_unify/modeling_qwen3.py
+++ b/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(),
     )
 
 
diff --git a/python/sglang/multimodal_gen/test/unit/test_sensenova_u1.py b/python/sglang/multimodal_gen/test/unit/test_sensenova_u1.py
index 1453f26eb..f96b9f0a4 100644
--- a/python/sglang/multimodal_gen/test/unit/test_sensenova_u1.py
+++ b/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()
diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py
index 5bb45333a..e03bd2632 100644
--- a/python/sglang/srt/layers/layernorm.py
+++ b/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.