From 628df31d088f612ffbc255b70b578cf9519d2a09 Mon Sep 17 00:00:00 2001
From: Thomas Wang <1am9trash@gmail.com>
Date: Thu, 9 Apr 2026 16:55:29 +0800
Subject: [PATCH] [AMD] Use aiter CK layernorm2d for LayerNorm to reduce NSA
 indexer kernel launches (#22424)

---
 .../srt/layers/attention/nsa/nsa_indexer.py      | 14 ++++++++++++--
 python/sglang/srt/layers/layernorm.py            | 16 +++++++++++++++-
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
index 6bfcb3f66..4ffd13bdd 100644
--- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
@@ -16,12 +16,20 @@ from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor
 from sglang.srt.layers.layernorm import LayerNorm
 from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.layers.utils import MultiPlatformOp
-from sglang.srt.utils import add_prefix, ceil_align, is_cuda, is_hip, is_npu
+from sglang.srt.utils import (
+    add_prefix,
+    ceil_align,
+    get_bool_env_var,
+    is_cuda,
+    is_hip,
+    is_npu,
+)
 
 global _use_multi_stream
 _is_cuda = is_cuda()
 _is_hip = is_hip()
 _is_npu = is_npu()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 _is_fp8_fnuz = is_fp8_fnuz()
 if _is_cuda:
     try:
@@ -212,7 +220,9 @@ class Indexer(MultiPlatformOp):
             params_dtype=torch.bfloat16 if _is_cuda else torch.float32,
             prefix=add_prefix("weights_proj", prefix),
         )
-        self.k_norm = LayerNorm(self.head_dim, dtype=torch.float32)
+        self.k_norm = LayerNorm(
+            self.head_dim, dtype=torch.bfloat16 if _use_aiter else torch.float32
+        )
         self.rotary_emb = get_rope_wrapper(
             rope_head_dim,
             rotary_dim=rope_head_dim,
diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py
index 0db6675e6..60eb942e1 100644
--- a/python/sglang/srt/layers/layernorm.py
+++ b/python/sglang/srt/layers/layernorm.py
@@ -65,11 +65,14 @@ if _is_cuda or _is_xpu:
         gemma_rmsnorm,
         rmsnorm,
     )
+_has_aiter_layer_norm = False
 _has_vllm_rms_norm = False
 if _use_aiter:
+    from aiter import layernorm2d_fwd as layer_norm
     from aiter import rmsnorm2d_fwd as rms_norm
     from aiter import rmsnorm2d_fwd_with_add as fused_add_rms_norm
 
+    _has_aiter_layer_norm = True  # aiter provides the layer_norm functions
     _has_vllm_rms_norm = True  # aiter provides the rms_norm functions
 elif _is_hip:
     try:
@@ -428,7 +431,18 @@ class LayerNorm(MultiPlatformOp):
         self,
         x: torch.Tensor,
     ) -> torch.Tensor:
-        return self.forward_native(x)
+        if (
+            _has_aiter_layer_norm
+            and x.dtype in (torch.bfloat16, torch.float16)
+            and x.dtype == self.dtype
+        ):
+            orig_shape = x.shape
+            x = x.reshape(-1, self.hidden_size)
+            return layer_norm(x, self.weight, self.bias, self.variance_epsilon).view(
+                orig_shape
+            )
+        else:
+            return self.forward_native(x)
 
     def forward_npu(
         self,