[Perf] [AMD] mla: Use the gfx950 aiter RoPE, RMSNorm and MLA KV-row paths on gfx1250

python/sglang/srt/layers/layernorm.py

@@ -40,7 +40,6 @@
     is_cpu,
     is_cuda,
     is_flashinfer_available,
-    is_gfx1250_supported,
     is_hip,
     is_musa,
     is_npu,
@@ -113,17 +112,8 @@ def layernorm(
 if _use_aiter:
     import aiter as _aiter
     from aiter import layernorm2d_fwd as layer_norm
-
-    if is_gfx1250_supported():
-        from aiter.ops.triton.normalization.rmsnorm import (
-            rms_norm,
-        )
-        from aiter.ops.triton.normalization.rmsnorm import (
-            rmsnorm2d_fwd_with_add as fused_add_rms_norm,
-        )
-    else:
-        from aiter import rmsnorm2d_fwd as rms_norm
-        from aiter import rmsnorm2d_fwd_with_add as fused_add_rms_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

python/sglang/srt/layers/rotary_embedding/factory.py

@@ -76,9 +76,10 @@ def _bailing_yarn_kwargs(rope_scaling: Dict[str, Any], max_position: int) -> Dic
 
 
 @functools.lru_cache(maxsize=1)
-def _aiter_rope_unsupported_arch() -> bool:
-    """aiter's rope kernels (csrc/kernels/rope/rope_common.h) depend on ck_tile
-    types that do not build on gfx1250; fall back to sglang's native rope there."""
+def _keeps_sglang_rope_modules() -> bool:
+    """gfx1250 keeps sglang's rope modules: its MLA paths read cos_sin_cache,
+    which aiter's rope modules do not have. DeepseekScalingRotaryEmbedding
+    calls the aiter rope kernel directly instead."""
     if not _is_hip:
         return False
     try:
@@ -525,7 +526,7 @@ def get_rope_wrapper(
     device: Optional[str] = None,
 ):
     if device != "cpu":
-        use_aiter_rope = _use_aiter and not _aiter_rope_unsupported_arch()
+        use_aiter_rope = _use_aiter and not _keeps_sglang_rope_modules()
         wrapper = aiter_get_rope if use_aiter_rope else get_rope
         return wrapper(
             head_size,

python/sglang/srt/layers/rotary_embedding/rope_variant.py

@@ -11,7 +11,7 @@
 import torch.nn.functional as F
 
 from sglang.kernels.fused_op import BaseFusedOp
-from sglang.srt.layers.rotary_embedding.base import RotaryEmbedding
+from sglang.srt.layers.rotary_embedding.base import RotaryEmbedding, _use_aiter
 from sglang.srt.layers.rotary_embedding.utils import (
     apply_rotary_pos_emb_native,
     rotate_gptj,
@@ -23,13 +23,24 @@
     yarn_get_mscale,
     yarn_linear_ramp_mask,
 )
-from sglang.srt.utils import cpu_has_amx_support, get_device, is_cuda, is_hip, is_npu
+from sglang.srt.utils import (
+    cpu_has_amx_support,
+    get_device,
+    is_cuda,
+    is_gfx1250_supported,
+    is_hip,
+    is_npu,
+)
 
 _is_cuda = is_cuda()
 _is_hip = is_hip()
 _is_npu = is_npu()
+_use_aiter_rope_gfx1250 = _use_aiter and is_gfx1250_supported()
 _is_cpu_amx_available = cpu_has_amx_support()
 
+if _use_aiter_rope_gfx1250:
+    import aiter
+
 if _is_npu:
     import torch_npu
 
@@ -374,9 +385,59 @@ def __init__(
         super().__init__(
             head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
         )
-        if _is_hip:
+        self._aiter_cos_sin = None
+        if _use_aiter_rope_gfx1250 and not self._force_native:
+            self._forward_method = self._forward_aiter_cached
+        elif _is_hip:
             self._forward_method = self.forward_native
 
+    def _build_aiter_cos_sin(self, dtype: torch.dtype) -> None:
+        # The kernel indexes cos/sin as contiguous [max_pos, rotary_dim / 2]
+        # without strides, so chunked views of cos_sin_cache cannot be passed.
+        cache = self.cos_sin_cache
+        cos, sin = cache.chunk(2, dim=-1)
+        shape = (cache.shape[0], 1, 1, cos.shape[-1])
+        self._aiter_cos_sin = (
+            cache,
+            cos.to(dtype).reshape(shape).contiguous(),
+            sin.to(dtype).reshape(shape).contiguous(),
+        )
+
+    def _forward_aiter_cached(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        offsets: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if offsets is not None:
+            # The cached-positions kernel takes no offsets.
+            return self.forward_native(positions, query, key, offsets)
+
+        cached = self._aiter_cos_sin
+        if (
+            cached is None
+            or cached[0] is not self.cos_sin_cache
+            or cached[1].dtype != query.dtype
+        ):
+            self._build_aiter_cos_sin(query.dtype)
+        _, cos, sin = self._aiter_cos_sin
+        num_tokens = positions.numel()
+        query_shape, key_shape = query.shape, key.shape
+        query = query.view(1, num_tokens, -1, self.head_size)
+        key = key.view(1, num_tokens, -1, self.head_size)
+        aiter.rope_cached_positions_2c_fwd_inplace(
+            query[..., : self.rotary_dim],
+            key[..., : self.rotary_dim],
+            cos,
+            sin,
+            positions.view(1, num_tokens),
+            0 if self.is_neox_style else 1,
+            reuse_freqs_front_part=True,
+            nope_first=False,
+        )
+        return query.view(query_shape), key.view(key_shape)
+
     def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
         pos_freqs = self.base ** (
             torch.arange(0, self.rotary_dim, 2, dtype=torch.float, device=self.device)

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py

@@ -32,6 +32,8 @@
 )
 from sglang.srt.models.deepseek_common.utils import (
     _is_block_scale_fp8,
+    _is_gfx1250_supported,
+    _use_aiter,
     _use_aiter_bpreshuffle_gfx95,
     _use_aiter_gfx95,
 )
@@ -44,6 +46,8 @@
 _use_fp8_prefill_attn = (
     get_bool_env_var("SGLANG_AITER_FP8_PREFILL_ATTN", "True") and _use_aiter_gfx95
 )
+# Same as gfx950: write/read MLA KV rows via store_mla/read_mla, not index_put.
+_use_mla_kv_rows = _use_aiter_gfx95 or (_use_aiter and _is_gfx1250_supported)
 
 if _use_aiter_gfx95:
     from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
@@ -329,7 +333,7 @@ def _set_mla_kv_buffer_rocm(
         k_pe: torch.Tensor,
         forward_batch: ForwardBatch,
     ):
-        if _use_aiter_gfx95:
+        if _use_mla_kv_rows:
             get_kv_layers().layer(self.attn_mha.layer_id).store_mla(
                 self.attn_mha,
                 KVWriteLoc.for_batch(forward_batch),
@@ -349,7 +353,7 @@ def _get_mla_kv_buffer_rocm(
         dst_dtype: torch.dtype,
         forward_batch: ForwardBatch,
     ):
-        if _use_aiter_gfx95:
+        if _use_mla_kv_rows:
             kv_a, k_pe = (
                 get_kv_layers()
                 .read(self.attn_mha.layer_id)

python/sglang/srt/models/deepseek_common/utils.py

@@ -50,10 +50,10 @@
 _is_xpu = is_xpu()
 _device_sm = get_device_sm()
 _is_gfx95_supported = is_gfx95_supported()
-# gfx1250 reuses the gfx95 (CDNA4) code paths for MXFP4 q/k-norm kernels, but its
-# aiter rope kernels (ck_tile) do not build, so it runs sglang's native rope which
-# lacks the separate cos_cache/sin_cache buffers the gfx95 fused-rope decode path
-# expects. This flag lets gfx1250 carve out of that fused-rope path.
+# gfx1250 reuses the gfx95 (CDNA4) code paths for MXFP4 q/k-norm kernels, but it
+# runs sglang's rope modules, which lack the separate cos_cache/sin_cache buffers
+# the gfx95 fused-rope decode path expects. This flag lets gfx1250 carve out of
+# that fused-rope path.
 _is_gfx1250_supported = is_gfx1250_supported()
 _use_aiter_gfx95 = _use_aiter and _is_gfx95_supported
 _use_aiter_bpreshuffle_gfx95 = _use_aiter_gfx95 and get_hip_version() >= (7, 2, 0)

test/registered/rotary/test_rope_rocm.py

@@ -1,8 +1,12 @@
 import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
 
 import torch
 
 from sglang.srt.layers.rotary_embedding import RotaryEmbedding
+from sglang.srt.layers.rotary_embedding import base as rotary_base
+from sglang.srt.layers.rotary_embedding import rope_variant
 from sglang.srt.utils import get_bool_env_var, is_hip
 from sglang.test.ci.ci_register import register_amd_ci
 from sglang.test.test_utils import CustomTestCase
@@ -140,5 +144,126 @@ def test_all_cases(self) -> None:
                 self._run_case_aiter(*case)
 
 
+@unittest.skipUnless(
+    rope_variant._use_aiter_rope_gfx1250,
+    "Requires gfx1250 plus SGLANG_USE_AITER=1",
+)
+class TestDeepseekScalingRopeAiterGfx1250(CustomTestCase):
+    # Shapes follow the MLA callers: q_pe is the rope slice of q, k_pe the rope
+    # slice of the latent KV row.
+    def _run_case(self, is_neox: bool, num_tokens: int, num_heads: int) -> None:
+        rope = rope_variant.DeepseekScalingRotaryEmbedding(
+            64,
+            64,
+            4096,
+            50000,
+            is_neox,
+            32.0,
+            torch.bfloat16,
+            mscale=1.0,
+            mscale_all_dim=1.0,
+            device="cuda",
+        ).to("cuda")
+        positions = torch.randint(0, 4096, (num_tokens,), device="cuda")
+        q = torch.randn(num_tokens, num_heads, 192, dtype=torch.bfloat16, device="cuda")
+        kv = torch.randn(num_tokens, 1, 576, dtype=torch.bfloat16, device="cuda")
+        q_pe, k_pe = q[..., 128:], kv[..., 512:]
+
+        q_ref, k_ref = rope.forward_native(
+            positions, q_pe.contiguous(), k_pe.contiguous()
+        )
+        q_out, k_out = rope(positions, q_pe, k_pe)
+
+        self.assertEqual(q_out.data_ptr(), q_pe.data_ptr())
+        self.assertEqual(k_out.data_ptr(), k_pe.data_ptr())
+        # forward_native adds a leading dim on the neox path.
+        torch.testing.assert_close(
+            q[..., 128:], q_ref.reshape(q_pe.shape), atol=3e-2, rtol=2e-2
+        )
+        torch.testing.assert_close(
+            kv[..., 512:], k_ref.reshape(k_pe.shape), atol=3e-2, rtol=2e-2
+        )
+
+    def test_matches_native(self) -> None:
+        for is_neox in (False, True):
+            for num_tokens, num_heads in ((1, 16), (311, 8), (2048, 64)):
+                with self.subTest(is_neox=is_neox, num_tokens=num_tokens):
+                    self._run_case(is_neox, num_tokens, num_heads)
+
+    def test_rebuilds_after_cache_growth(self) -> None:
+        rope = self._small_rope()
+        q = torch.randn(4, 16, 64, dtype=torch.bfloat16, device="cuda")
+        k = torch.randn(4, 1, 64, dtype=torch.bfloat16, device="cuda")
+        rope(torch.arange(4, device="cuda"), q.clone(), k.clone())
+
+        old_cache = rope.cos_sin_cache
+        rope._ensure_cos_sin_cache_length(1024)
+        self.assertIsNot(rope.cos_sin_cache, old_cache)
+        positions = torch.tensor([0, 255, 600, 1000], device="cuda")
+        q_ref, k_ref = rope.forward_native(positions, q.clone(), k.clone())
+        q_out, k_out = rope(positions, q.clone(), k.clone())
+        self.assertIs(rope._aiter_cos_sin[0], rope.cos_sin_cache)
+        torch.testing.assert_close(q_out, q_ref, atol=3e-2, rtol=2e-2)
+        torch.testing.assert_close(k_out, k_ref, atol=3e-2, rtol=2e-2)
+
+    def _small_rope(self) -> rope_variant.DeepseekScalingRotaryEmbedding:
+        return rope_variant.DeepseekScalingRotaryEmbedding(
+            64,
+            64,
+            256,
+            50000,
+            False,
+            1.0,
+            torch.bfloat16,
+            mscale=1.0,
+            mscale_all_dim=1.0,
+            device="cuda",
+        ).to("cuda")
+
+    def test_rebuilds_on_dtype_change(self) -> None:
+        rope = self._small_rope()
+        positions = torch.arange(4, device="cuda")
+        for dtype in (torch.bfloat16, torch.float16, torch.bfloat16):
+            q = torch.randn(4, 16, 64, dtype=dtype, device="cuda")
+            k = torch.randn(4, 1, 64, dtype=dtype, device="cuda")
+            q_ref, k_ref = rope.forward_native(positions, q.clone(), k.clone())
+            q_out, k_out = rope(positions, q.clone(), k.clone())
+            self.assertEqual(rope._aiter_cos_sin[1].dtype, dtype)
+            torch.testing.assert_close(q_out, q_ref, atol=3e-2, rtol=2e-2)
+            torch.testing.assert_close(k_out, k_ref, atol=3e-2, rtol=2e-2)
+
+    def test_offsets_use_native(self) -> None:
+        rope = self._small_rope()
+        positions = torch.arange(4, device="cuda")
+        q = torch.randn(4, 16, 64, dtype=torch.bfloat16, device="cuda")
+        k = torch.randn(4, 1, 64, dtype=torch.bfloat16, device="cuda")
+        offsets = torch.full((4,), 7, device="cuda")
+        q_ref, k_ref = rope.forward_native(positions, q.clone(), k.clone(), offsets)
+        q_out, k_out = rope(positions, q.clone(), k.clone(), offsets)
+        torch.testing.assert_close(q_out, q_ref, atol=3e-2, rtol=2e-2)
+        torch.testing.assert_close(k_out, k_ref, atol=3e-2, rtol=2e-2)
+
+    def test_key_passed_as_its_own_query(self) -> None:
+        # IndexerKPool's K-only path passes the rope slice of the key twice.
+        rope = self._small_rope()
+        positions = torch.arange(4, device="cuda")
+        key = torch.randn(4, 1, 128, dtype=torch.bfloat16, device="cuda")
+        k_rope = key[..., :64]
+        k_ref = rope.forward_native(positions, k_rope.clone(), k_rope.clone())[1]
+        rope(positions, k_rope, k_rope)
+        torch.testing.assert_close(key[..., :64], k_ref, atol=3e-2, rtol=2e-2)
+
+    def test_force_native_keeps_native(self) -> None:
+        exec_ctx = SimpleNamespace(
+            deterministic=SimpleNamespace(rl_on_policy_target="fsdp")
+        )
+        with (
+            patch.object(rotary_base, "publish_role", return_value="rollout"),
+            patch.object(rotary_base, "get_exec", return_value=exec_ctx),
+        ):
+            rope = self._small_rope()
+        self.assertEqual(rope._forward_method, rope.forward_native)
+
+
 if __name__ == "__main__":
     unittest.main()