[CPU] refactor rope kernels (#31897)

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

sgl-kernel/csrc/cpu/vec.h

@@ -44,6 +44,17 @@ inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const fl
   return std::make_tuple(x0, x1);
 }
 
+template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 1>
+inline at::vec::Vectorized<float> load_float_vec(const scalar_t* __restrict__ data) {
+  at::vec::Vectorized<float> out;
+  if constexpr (std::is_same_v<scalar_t, at::BFloat16>) {
+    at::vec::load_fp32_from_bf16(data, out);
+  } else {
+    at::vec::load_fp32_from_fp16(data, out);
+  }
+  return out;
+}
+
 #if defined(CPU_CAPABILITY_AVX512)
 
 // `at::vec::convert_from_float<>` from PyTorch doesn't have avx512-bf16 intrinsics

test/registered/cpu/test_rope.py

@@ -161,7 +161,6 @@ class TestROPE(CustomTestCase):
                 atol = rtol = precision[q_pe.dtype]
                 torch.testing.assert_close(q_pe, q_pe_clone, atol=atol, rtol=rtol)
                 torch.testing.assert_close(k_pe, k_pe_clone, atol=atol, rtol=rtol)
-                torch.testing.assert_close(k_pe, k_pe_clone)
 
     def test_origin_rope(self):
         def single_test(