[AMD] Replace naive triton RMSNorm with aiter RMSNorm for diffusion model (#24360)

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -28,10 +28,12 @@ from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
 
 _is_cuda = current_platform.is_cuda()
+_is_hip = current_platform.is_hip()
 _is_npu = current_platform.is_npu()
 _is_musa = current_platform.is_musa()
 _is_cpu = current_platform.is_cpu()
 _is_xpu = current_platform.is_xpu()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 
 if _is_cuda or _is_xpu:
     from sgl_kernel import fused_add_rmsnorm, rmsnorm
@@ -41,6 +43,11 @@ if _is_npu:
 
 if _is_musa:
     from sgl_kernel import fused_add_rmsnorm
+
+if _use_aiter:
+    from aiter import rmsnorm2d_fwd as rms_norm
+    from aiter import rmsnorm2d_fwd_with_add as fused_add_rms_norm
+
 if not _is_cpu:
     from sglang.jit_kernel.diffusion.triton.norm import norm_infer, rms_norm_fn
 
@@ -70,6 +77,8 @@ class RMSNorm(CustomOp):
         )
         if get_bool_env_var("SGLANG_ENABLE_DETERMINISTIC_INFERENCE"):
             self._forward_method = self.forward_native
+        elif _use_aiter:
+            self._forward_method = self.forward_aiter
 
     def forward_triton(self, x: torch.Tensor, residual: Optional[torch.Tensor] = None):
         return rms_norm_fn(
@@ -177,6 +186,47 @@ class RMSNorm(CustomOp):
         # ROCm builds of sgl-kernel do not expose rmsnorm custom ops yet.
         return self.forward_native(x, residual)
 
+    def forward_aiter(
+        self,
+        x: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        # Fall back to the native fp32 path for cases aiter cannot serve:
+        #   - fp32 input  (CK kernel is templated on fp16/bf16 only;
+        #                  out.dtype check rejects fp32 with "not support output type: float")
+        if (
+            x.dtype not in (torch.float16, torch.bfloat16)
+            or self.variance_size_override is not None
+        ):
+            return self.forward_native(x, residual)
+
+        weight = self._get_weight(x.dtype)
+
+        shape = x.shape
+        x_2d = x.reshape(
+            -1, shape[-1]
+        )  # (bs, seq_len, hidden_size) -> (bs*seq_len, hidden_size)
+        if not x_2d.is_contiguous():
+            x_2d = x_2d.contiguous()
+
+        if residual is not None:
+            residual_shape = residual.shape
+            residual_2d = residual.reshape(-1, shape[-1])
+            if not residual_2d.is_contiguous():
+                residual_2d = residual_2d.contiguous()
+            output = torch.empty_like(x_2d)
+            residual_out = torch.empty_like(x_2d)
+            fused_add_rms_norm(
+                output,
+                x_2d,
+                residual_2d,
+                residual_out,
+                weight,
+                self.variance_epsilon,
+            )
+            return output.view(shape), residual_out.view(residual_shape)
+        return rms_norm(x_2d, weight, self.variance_epsilon).view(shape)
+
     def _get_weight(self, dtype: torch.dtype) -> torch.Tensor:
         """Return weight matched to *dtype*.