From 10e0b83a4c65643ba16c99a5e55d54e9a4b365b4 Mon Sep 17 00:00:00 2001
From: iLeGend <824040212@qq.com>
Date: Thu, 20 Nov 2025 13:19:53 +0800
Subject: [PATCH] Add FP32 dtype support for RoPE - Part2 (#13328)

---
 python/sglang/srt/layers/rotary_embedding.py | 8 ++------
 test/srt/cpu/test_rope.py                    | 6 ++++++
 test/srt/rotary_embedding/test_mrope.py      | 2 +-
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py
index 5f62b3f2e..dfddc1d67 100644
--- a/python/sglang/srt/layers/rotary_embedding.py
+++ b/python/sglang/srt/layers/rotary_embedding.py
@@ -113,7 +113,7 @@ class RotaryEmbedding(CustomOp):
         if not _is_cuda:
             cache = cache.to(dtype)
 
-        if dtype == torch.float32 or (
+        if (
             (not (_is_cuda or _is_npu) or self.head_size not in [64, 128, 256, 512])
             and not (_is_cpu and _is_cpu_amx_available)
             and not (_is_xpu)
@@ -273,11 +273,7 @@ class RotaryEmbedding(CustomOp):
         offsets: Optional[torch.Tensor] = None,
         fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        if (
-            _is_cuda
-            and (self.head_size in [64, 128, 256, 512])
-            and self.dtype != torch.float32
-        ):
+        if _is_cuda and (self.head_size in [64, 128, 256, 512]):
             apply_rope_with_cos_sin_cache_inplace(
                 positions=positions,
                 query=query,
diff --git a/test/srt/cpu/test_rope.py b/test/srt/cpu/test_rope.py
index 8c1dfe9aa..22824e0ca 100644
--- a/test/srt/cpu/test_rope.py
+++ b/test/srt/cpu/test_rope.py
@@ -146,6 +146,12 @@ class TestROPE(CustomTestCase):
             (128, 128, 2048, 10000, False, torch.bfloat16, "cpu", 2, 512, 32, 8),
             (128, 128, 2048, 10000, False, torch.bfloat16, "cpu", 2, 512, 16, 4),
             (512, 128, 311, 10000, False, torch.bfloat16, "cpu", 3, 39, 4, 2),
+            (64, 64, 32, 8000, True, torch.float32, "cpu", 32, 32, 1, 1),
+            (256, 128, 4096, 10000, True, torch.float32, "cpu", 2, 512, 32, 8),
+            (512, 128, 311, 10000, True, torch.float32, "cpu", 3, 39, 4, 2),
+            (128, 128, 2048, 10000, False, torch.float32, "cpu", 2, 512, 32, 8),
+            (128, 128, 2048, 10000, False, torch.float32, "cpu", 2, 512, 16, 4),
+            (512, 128, 311, 10000, False, torch.float32, "cpu", 3, 39, 4, 2),
         ]
 
         for (
diff --git a/test/srt/rotary_embedding/test_mrope.py b/test/srt/rotary_embedding/test_mrope.py
index ad6412ec4..4fbfd0691 100644
--- a/test/srt/rotary_embedding/test_mrope.py
+++ b/test/srt/rotary_embedding/test_mrope.py
@@ -76,7 +76,7 @@ num_tokens_list = [11, 8192]
     ],
 )
 @pytest.mark.parametrize("tp_size", [1, 2])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
 @pytest.mark.parametrize("num_tokens", num_tokens_list)
 def test_mrope(
     model_name: str,