[CPU] Implement fused QK Norm and RoPE kernels (#37748)

Signed-off-by: Cui, Lily <lily.cui@intel.com>

python/sglang/kernels/aot/csrc/cpu/norm.cpp

@@ -1052,3 +1052,280 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_qk_gemma_rmsnorm_with_gate_
   });
   return std::make_tuple(q_out, k_out, gate_out);
 }
+
+namespace {
+
+template <typename scalar_t>
+inline void
+fused_qk_norm_per_head(scalar_t* __restrict__ data, const scalar_t* __restrict__ weight, int64_t D, float eps) {
+  using bVec = at::vec::Vectorized<scalar_t>;
+  using fVec = at::vec::Vectorized<float>;
+  constexpr int64_t kVecSize = bVec::size();
+
+  fVec sum2_fvec{0.f};
+  float sum2_val{0.f};
+
+  int64_t d = 0;
+#pragma GCC unroll 4
+  for (; d <= D - kVecSize; d += kVecSize) {
+    auto [x_fvec0, x_fvec1] = load_float_vec2(data + d);
+    sum2_fvec += x_fvec0 * x_fvec0;
+    sum2_fvec += x_fvec1 * x_fvec1;
+  }
+  for (; d < D; ++d) {
+    const float x_val = static_cast<float>(data[d]);
+    sum2_val += x_val * x_val;
+  }
+
+  const float scale = 1.f / std::sqrt((sum2_val + vec_reduce_sum(sum2_fvec)) / D + eps);
+  const fVec scale_fvec{scale};
+
+  d = 0;
+#pragma GCC unroll 4
+  for (; d <= D - kVecSize; d += kVecSize) {
+    auto [x_fvec0, x_fvec1] = load_float_vec2(data + d);
+    auto [w_fvec0, w_fvec1] = load_float_vec2(weight + d);
+    convert_from_float_ext<scalar_t>(x_fvec0 * scale_fvec * w_fvec0, x_fvec1 * scale_fvec * w_fvec1).store(data + d);
+  }
+  for (; d < D; ++d) {
+    data[d] = static_cast<scalar_t>(static_cast<float>(data[d]) * scale * static_cast<float>(weight[d]));
+  }
+}
+
+template <typename scalar_t>
+void fused_qk_norm_kernel_impl(
+    scalar_t* __restrict__ q,
+    scalar_t* __restrict__ k,
+    const scalar_t* __restrict__ q_weight,
+    const scalar_t* __restrict__ k_weight,
+    int64_t num_tokens,
+    int64_t num_q_heads,
+    int64_t num_kv_heads,
+    int64_t head_dim,
+    int64_t q_stride,
+    int64_t k_stride,
+    float eps) {
+  const int64_t num_qk_heads = num_q_heads + num_kv_heads;
+
+  at::parallel_for(0, num_tokens * num_qk_heads, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t work = begin; work < end; ++work) {
+      const int64_t token = work / num_qk_heads;
+      const int64_t local_head = work % num_qk_heads;
+      const bool is_q = local_head < num_q_heads;
+
+      scalar_t* __restrict__ data = is_q ? q + token * q_stride + local_head * head_dim
+                                         : k + token * k_stride + (local_head - num_q_heads) * head_dim;
+      fused_qk_norm_per_head<scalar_t>(data, is_q ? q_weight : k_weight, head_dim, eps);
+    }
+  });
+}
+
+template <typename scalar_t>
+inline void fused_qk_norm_rope_apply_interleaved(
+    scalar_t* __restrict__ data, const scalar_t* __restrict__ cache, int64_t rotary_dim) {
+  constexpr int64_t kVecSize = at::vec::Vectorized<scalar_t>::size();
+  const int64_t half_rotary = rotary_dim / 2;
+
+  int64_t d = 0;
+  for (; d <= rotary_dim - kVecSize; d += kVecSize) {
+    auto [xy0, xy1] = load_float_vec2(data + d);
+    auto [x, y] = at::vec::deinterleave2(xy0, xy1);
+    auto cos = load_float_vec(cache + d / 2);
+    auto sin = load_float_vec(cache + half_rotary + d / 2);
+    auto out0 = x * cos - y * sin;
+    auto out1 = y * cos + x * sin;
+    std::tie(xy0, xy1) = at::vec::interleave2(out0, out1);
+    convert_from_float_ext<scalar_t>(xy0, xy1).store(data + d);
+  }
+  for (; d < rotary_dim; d += 2) {
+    const float x = static_cast<float>(data[d]);
+    const float y = static_cast<float>(data[d + 1]);
+    const float c = static_cast<float>(cache[d / 2]);
+    const float s = static_cast<float>(cache[half_rotary + d / 2]);
+    data[d] = static_cast<scalar_t>(x * c - y * s);
+    data[d + 1] = static_cast<scalar_t>(y * c + x * s);
+  }
+}
+
+template <typename scalar_t>
+inline void
+fused_qk_norm_rope_apply_neox(scalar_t* __restrict__ data, const scalar_t* __restrict__ cache, int64_t rotary_dim) {
+  constexpr int64_t kVecSize = at::vec::Vectorized<scalar_t>::size();
+  const int64_t half_rotary = rotary_dim / 2;
+
+  int64_t d = 0;
+  for (; d <= half_rotary - kVecSize; d += kVecSize) {
+    auto [x0, x1] = load_float_vec2(data + d);
+    auto [y0, y1] = load_float_vec2(data + half_rotary + d);
+    auto [cos0, cos1] = load_float_vec2(cache + d);
+    auto [sin0, sin1] = load_float_vec2(cache + half_rotary + d);
+    auto out0 = x0 * cos0 - y0 * sin0;
+    auto out1 = x1 * cos1 - y1 * sin1;
+    auto out2 = y0 * cos0 + x0 * sin0;
+    auto out3 = y1 * cos1 + x1 * sin1;
+    convert_from_float_ext<scalar_t>(out0, out1).store(data + d);
+    convert_from_float_ext<scalar_t>(out2, out3).store(data + half_rotary + d);
+  }
+  for (; d < half_rotary; ++d) {
+    const float x = static_cast<float>(data[d]);
+    const float y = static_cast<float>(data[d + half_rotary]);
+    const float c = static_cast<float>(cache[d]);
+    const float s = static_cast<float>(cache[half_rotary + d]);
+    data[d] = static_cast<scalar_t>(x * c - y * s);
+    data[d + half_rotary] = static_cast<scalar_t>(y * c + x * s);
+  }
+}
+
+template <typename scalar_t>
+inline void fused_qk_norm_rope_per_head(
+    scalar_t* __restrict__ data,
+    const scalar_t* __restrict__ weight,
+    int64_t head_dim,
+    int64_t rotary_dim,
+    const scalar_t* __restrict__ cache_row,
+    bool is_neox,
+    float eps) {
+  fused_qk_norm_per_head<scalar_t>(data, weight, head_dim, eps);
+
+  if (is_neox) {
+    fused_qk_norm_rope_apply_neox<scalar_t>(data, cache_row, rotary_dim);
+  } else {
+    fused_qk_norm_rope_apply_interleaved<scalar_t>(data, cache_row, rotary_dim);
+  }
+}
+
+template <typename scalar_t>
+void fused_qk_norm_rope_kernel_impl(
+    scalar_t* __restrict__ q,
+    scalar_t* __restrict__ k,
+    const scalar_t* __restrict__ q_weight,
+    const scalar_t* __restrict__ k_weight,
+    int64_t num_tokens,
+    int64_t num_q_heads,
+    int64_t num_kv_heads,
+    int64_t head_dim,
+    int64_t q_stride,
+    int64_t k_stride,
+    float eps,
+    bool is_neox,
+    const int64_t* __restrict__ position_ids,
+    const scalar_t* __restrict__ cos_sin_cache,
+    int64_t rotary_dim) {
+  const int64_t num_qk_heads = num_q_heads + num_kv_heads;
+  at::parallel_for(0, num_tokens * num_qk_heads, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t work = begin; work < end; ++work) {
+      const int64_t token = work / num_qk_heads;
+      const int64_t local_head = work % num_qk_heads;
+      const bool is_q = local_head < num_q_heads;
+
+      scalar_t* __restrict__ data = is_q ? q + token * q_stride + local_head * head_dim
+                                         : k + token * k_stride + (local_head - num_q_heads) * head_dim;
+      const scalar_t* __restrict__ cache_row = cos_sin_cache + position_ids[token] * rotary_dim;
+      fused_qk_norm_rope_per_head<scalar_t>(
+          data, is_q ? q_weight : k_weight, head_dim, rotary_dim, cache_row, is_neox, eps);
+    }
+  });
+}
+
+}  // anonymous namespace
+
+void fused_qk_norm_cpu(
+    at::Tensor& q, at::Tensor& k, const at::Tensor& q_weight, const at::Tensor& k_weight, double eps) {
+  const auto st = q.scalar_type();
+  CHECK_INPUT_ND<2>(q);
+  CHECK_INPUT_ND<2>(k);
+  CHECK_EQ(k.size(0), q.size(0));
+  CHECK_EQ(k.scalar_type(), st);
+
+  const int64_t head_dim = q_weight.numel();
+  CHECK_GT(head_dim, 0);
+  CHECK_INPUT_SHAPE_DTYPE<false>(q_weight, {head_dim}, st);
+  CHECK_INPUT_SHAPE_DTYPE<false>(k_weight, {head_dim}, st);
+  CHECK_EQ(q.size(1) % head_dim, 0);
+  CHECK_EQ(k.size(1) % head_dim, 0);
+
+  const int64_t num_tokens = q.size(0);
+  if (num_tokens == 0) return;
+
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_qk_norm_kernel", [&] {
+    fused_qk_norm_kernel_impl<scalar_t>(
+        q.data_ptr<scalar_t>(),
+        k.data_ptr<scalar_t>(),
+        q_weight.data_ptr<scalar_t>(),
+        k_weight.data_ptr<scalar_t>(),
+        num_tokens,
+        q.size(1) / head_dim,
+        k.size(1) / head_dim,
+        head_dim,
+        q.stride(0),
+        k.stride(0),
+        static_cast<float>(eps));
+  });
+}
+
+void fused_qk_norm_rope_cpu(
+    at::Tensor& q,
+    at::Tensor& k,
+    const at::Tensor& q_weight,
+    const at::Tensor& k_weight,
+    double eps,
+    bool is_neox,
+    const at::Tensor& position_ids,
+    const at::Tensor& cos_sin_cache,
+    int64_t rotary_dim) {
+  const auto st = q.scalar_type();
+  CHECK_INPUT_ND<2>(q);
+  CHECK_INPUT_ND<2>(k);
+  CHECK_EQ(k.size(0), q.size(0));
+  CHECK_EQ(k.scalar_type(), st);
+  CHECK_DIM(1, position_ids);
+  CHECK_EQ(position_ids.size(0), q.size(0));
+  TORCH_CHECK(
+      position_ids.scalar_type() == at::kLong || position_ids.scalar_type() == at::kInt,
+      "position_ids must be int32 or int64, got ",
+      position_ids.scalar_type());
+  CHECK_INPUT_ND<2>(cos_sin_cache);
+  CHECK_EQ(cos_sin_cache.scalar_type(), st);
+  CHECK_EQ(cos_sin_cache.size(1), rotary_dim);
+
+  const int64_t head_dim = q_weight.numel();
+  CHECK_GT(head_dim, 0);
+  CHECK_INPUT_SHAPE_DTYPE<false>(q_weight, {head_dim}, st);
+  CHECK_INPUT_SHAPE_DTYPE<false>(k_weight, {head_dim}, st);
+  CHECK_EQ(q.size(1) % head_dim, 0);
+  CHECK_EQ(k.size(1) % head_dim, 0);
+  TORCH_CHECK(rotary_dim > 0 && rotary_dim <= head_dim, "rotary_dim must be in (0, head_dim]");
+  TORCH_CHECK(rotary_dim % 2 == 0, "rotary_dim must be even, got ", rotary_dim);
+
+  const int64_t num_tokens = q.size(0);
+  if (num_tokens == 0) return;
+
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_qk_norm_rope_kernel", [&] {
+    std::vector<int64_t> position_ids_i64;
+    const int64_t* pos_ptr;
+    if (position_ids.scalar_type() == at::kInt) {
+      position_ids_i64.resize(num_tokens);
+      const int* position_ids_i32 = position_ids.data_ptr<int>();
+      std::copy(position_ids_i32, position_ids_i32 + num_tokens, position_ids_i64.begin());
+      pos_ptr = position_ids_i64.data();
+    } else {
+      pos_ptr = position_ids.data_ptr<int64_t>();
+    }
+    fused_qk_norm_rope_kernel_impl<scalar_t>(
+        q.data_ptr<scalar_t>(),
+        k.data_ptr<scalar_t>(),
+        q_weight.data_ptr<scalar_t>(),
+        k_weight.data_ptr<scalar_t>(),
+        num_tokens,
+        q.size(1) / head_dim,
+        k.size(1) / head_dim,
+        head_dim,
+        q.stride(0),
+        k.stride(0),
+        static_cast<float>(eps),
+        is_neox,
+        pos_ptr,
+        cos_sin_cache.data_ptr<scalar_t>(),
+        rotary_dim);
+  });
+}

python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp

@@ -58,6 +58,20 @@ at::Tensor fused_add_layernorm_cpu(
     const std::optional<at::Tensor>& bias,
     double eps);
 
+// fused_qk_norm (per-head, in place)
+void fused_qk_norm_cpu(
+    at::Tensor& q, at::Tensor& k, const at::Tensor& q_weight, const at::Tensor& k_weight, double eps);
+void fused_qk_norm_rope_cpu(
+    at::Tensor& q,
+    at::Tensor& k,
+    const at::Tensor& q_weight,
+    const at::Tensor& k_weight,
+    double eps,
+    bool is_neox,
+    const at::Tensor& position_ids,
+    const at::Tensor& cos_sin_cache,
+    int64_t rotary_dim);
+
 // fused_qk_rmsnorm
 std::tuple<at::Tensor, at::Tensor> fused_qk_rmsnorm_cpu(
     const at::Tensor& q, const at::Tensor& k, const at::Tensor& q_weight, const at::Tensor& k_weight, double eps);
@@ -612,6 +626,12 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "fused_add_layernorm_cpu(Tensor input, Tensor residual, Tensor weight, Tensor? bias, float eps) -> "
       "Tensor");
   m.impl("fused_add_layernorm_cpu", torch::kCPU, &fused_add_layernorm_cpu);
+  m.def("fused_qk_norm_cpu(Tensor(a!) q, Tensor(b!) k, Tensor q_weight, Tensor k_weight, float eps) -> ()");
+  m.impl("fused_qk_norm_cpu", torch::kCPU, &fused_qk_norm_cpu);
+  m.def(
+      "fused_qk_norm_rope_cpu(Tensor(a!) q, Tensor(b!) k, Tensor q_weight, Tensor k_weight, float eps, "
+      "bool is_neox, Tensor position_ids, Tensor cos_sin_cache, int rotary_dim) -> ()");
+  m.impl("fused_qk_norm_rope_cpu", torch::kCPU, &fused_qk_norm_rope_cpu);
   m.def(
       "fused_qk_rmsnorm_cpu(Tensor q, Tensor k, Tensor q_weight, Tensor k_weight, float eps) -> "
       "(Tensor, Tensor)");

python/sglang/srt/models/qwen3_moe.py

@@ -19,6 +19,7 @@
 
 import logging
 import math
+from functools import lru_cache
 from typing import Any, Dict, Iterable, List, Optional, Tuple, TypeVar
 
 import torch
@@ -71,6 +72,7 @@ from sglang.srt.runtime_context import get_exec, get_forward, get_parallel, get_
 from sglang.srt.utils import (
     LazyValue,
     add_prefix,
+    is_cpu,
     is_cuda,
     is_flashinfer_available,
     is_non_idle_and_non_empty,
@@ -79,6 +81,7 @@ from sglang.srt.utils import (
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
 _is_cuda = is_cuda()
+_is_cpu = is_cpu()
 
 if _is_cuda:
     from sglang.kernels.ops.attention.fused_qknorm_rope import (
@@ -86,6 +89,12 @@ if _is_cuda:
         fused_qk_norm_rope,
     )
 
+
+@lru_cache(maxsize=1)
+def _has_cpu_fused_qk_norm_rope() -> bool:
+    return hasattr(torch.ops.sgl_kernel, "fused_qk_norm_rope_cpu")
+
+
 TConfig = TypeVar("TConfig", bound=PretrainedConfig)
 
 Qwen3MoeConfig = None
@@ -527,6 +536,12 @@ class Qwen3MoeAttention(nn.Module):
                 _yarn_factor != 1.0,
             )
         )
+        self.use_fused_qk_norm_rope_cpu = (
+            _is_cpu
+            and not isinstance(self.rotary_emb, MRotaryEmbedding)
+            and self.rotary_emb.rotary_dim % 2 == 0
+            and _has_cpu_fused_qk_norm_rope()
+        )
         self._used_fused_qk_norm_rope_last_call = False
 
         self.attn = RadixAttention(
@@ -594,31 +609,53 @@ class Qwen3MoeAttention(nn.Module):
         return None, forward_batch, inner_state
 
     def apply_qk_norm_rope(self, qkv, positions, forward_batch):
-        use_fused = self.use_fused_qk_norm_rope and qkv.dtype == torch.bfloat16
+        use_fused = (self.use_fused_qk_norm_rope and qkv.dtype == torch.bfloat16) or (
+            self.use_fused_qk_norm_rope_cpu
+            and qkv.dtype in (torch.bfloat16, torch.float16)
+        )
         if use_fused:
-            theta = self.rope_theta
-            positions = (
-                positions.view(-1).to(dtype=torch.int32, device=qkv.device).contiguous()
-            )
-            factor, low, high, attention_factor = compute_yarn_parameters(self.config)
-            fused_qk_norm_rope(
-                qkv,
-                self.num_heads,
-                self.num_kv_heads,
-                self.num_kv_heads,
-                self.head_dim,
-                self.q_norm.variance_epsilon,
-                self.q_norm.weight,
-                self.k_norm.weight,
-                theta,
-                self.rotary_emb.is_neox_style,
-                positions,
-                factor,
-                low,
-                high,
-                attention_factor,
-            )
-            q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+            if _is_cuda:
+                theta = self.rope_theta
+                positions = (
+                    positions.view(-1)
+                    .to(dtype=torch.int32, device=qkv.device)
+                    .contiguous()
+                )
+                factor, low, high, attention_factor = compute_yarn_parameters(
+                    self.config
+                )
+                fused_qk_norm_rope(
+                    qkv,
+                    self.num_heads,
+                    self.num_kv_heads,
+                    self.num_kv_heads,
+                    self.head_dim,
+                    self.q_norm.variance_epsilon,
+                    self.q_norm.weight,
+                    self.k_norm.weight,
+                    theta,
+                    self.rotary_emb.is_neox_style,
+                    positions,
+                    factor,
+                    low,
+                    high,
+                    attention_factor,
+                )
+                q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+            elif _is_cpu:
+                q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+                self.rotary_emb._match_cos_sin_cache_dtype(q)
+                torch.ops.sgl_kernel.fused_qk_norm_rope_cpu(
+                    q,
+                    k,
+                    self.q_norm.weight,
+                    self.k_norm.weight,
+                    self.q_norm.variance_epsilon,
+                    self.rotary_emb.is_neox_style,
+                    positions.view(-1),
+                    self.rotary_emb.cos_sin_cache,
+                    self.rotary_emb.rotary_dim,
+                )
             self._used_fused_qk_norm_rope_last_call = True
         else:
             # Fallback to non-fused QK Norm & RoPE implementation

python/sglang/srt/models/utils.py

@@ -38,7 +38,7 @@ from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
 from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.runtime_context import get_exec
-from sglang.srt.utils import get_current_device_stream_fast, is_cuda, is_hip
+from sglang.srt.utils import get_current_device_stream_fast, is_cpu, is_cuda, is_hip
 from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:
@@ -46,6 +46,7 @@ if TYPE_CHECKING:
 
 _is_cuda = is_cuda()
 _is_hip = is_hip()
+_is_cpu = is_cpu()
 
 WeightsMapping = Mapping[str, Optional[str]]
 """If a key maps to a value of `None`, the corresponding weight is ignored."""
@@ -449,6 +450,30 @@ def _reshape_for_qk_norm(x: torch.Tensor, head_dim: int) -> torch.Tensor:
     return x.reshape(-1, head_dim)
 
 
+@lru_cache(maxsize=1)
+def _has_cpu_fused_qk_norm() -> bool:
+    return hasattr(torch.ops.sgl_kernel, "fused_qk_norm_cpu")
+
+
+def can_use_fused_qk_norm_cpu(
+    q: torch.Tensor, k: torch.Tensor, head_dim: int, q_eps: float, k_eps: float
+) -> bool:
+    return (
+        _is_cpu
+        and q_eps == k_eps
+        and q.dim() == 2
+        and k.dim() == 2
+        and q.dtype in (torch.bfloat16, torch.float16)
+        and k.dtype == q.dtype
+        # q/k are usually strided views into qkv; only the head rows must be dense
+        and q.stride(-1) == 1
+        and k.stride(-1) == 1
+        and q.size(-1) % head_dim == 0
+        and k.size(-1) % head_dim == 0
+        and _has_cpu_fused_qk_norm()
+    )
+
+
 def apply_qk_norm(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -479,6 +504,12 @@ def apply_qk_norm(
     q_eps = q_norm.variance_epsilon
     k_eps = k_norm.variance_epsilon
 
+    if allow_inplace and can_use_fused_qk_norm_cpu(q, k, head_dim, q_eps, k_eps):
+        torch.ops.sgl_kernel.fused_qk_norm_cpu(
+            q, k, q_norm.weight, k_norm.weight, q_eps
+        )
+        return q, k
+
     if (
         _is_cuda  # TODO(dark): have not tested on ROCm or other backends
         and allow_inplace  # TODO(dark): this can be relaxed if needed

test/registered/cpu/test_norm.py

@@ -471,5 +471,132 @@ class TestFusedQKGemmaRMSNorm:
         )
 
 
+class TestFusedQKNorm:
+    def _norm_per_head_native(
+        self, x: torch.Tensor, weight: torch.Tensor, head_dim: int, eps: float
+    ):
+        out = x.reshape(-1, head_dim).float()
+        out = out * torch.rsqrt(out.pow(2).mean(-1, keepdim=True) + eps)
+        out = out.to(x.dtype) * weight
+        return out.reshape(x.shape)
+
+    @pytest.mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS)
+    @pytest.mark.parametrize(
+        "batch_size,num_head,num_head_kv,head_dim",
+        [
+            (1, 16, 2, 128),
+            (9, 8, 1, 64),
+            (256, 16, 2, 128),
+            (4109, 3, 1, 96),
+        ],
+    )
+    def test_fused_qk_norm(
+        self, batch_size: int, num_head: int, num_head_kv: int, head_dim: int, dtype
+    ):
+        q_size = num_head * head_dim
+        kv_size = num_head_kv * head_dim
+        qkv = torch.randn([batch_size, q_size + 2 * kv_size], dtype=dtype)
+        q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
+        v_before = v.clone()
+
+        q_weight = torch.randn(head_dim, dtype=dtype)
+        k_weight = torch.randn(head_dim, dtype=dtype)
+
+        ref_q_out = self._norm_per_head_native(q, q_weight, head_dim, eps)
+        ref_k_out = self._norm_per_head_native(k, k_weight, head_dim, eps)
+
+        torch.ops.sgl_kernel.fused_qk_norm_cpu(q, k, q_weight, k_weight, eps)
+
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(q, ref_q_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(k, ref_k_out, atol=atol, rtol=rtol)
+        # The in-place write must not spill past the q/k head rows into V.
+        torch.testing.assert_close(v, v_before, atol=0, rtol=0)
+
+    @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bfloat16"])
+    @pytest.mark.parametrize("is_neox", [False, True], ids=["interleaved", "neox"])
+    @pytest.mark.parametrize(
+        "num_head,num_head_kv,head_dim", [(16, 2, 64), (8, 1, 128)]
+    )
+    def test_fused_qk_norm_rope(
+        self,
+        dtype,
+        is_neox: bool,
+        num_head: int,
+        num_head_kv: int,
+        head_dim: int,
+    ):
+        batch_size = 3
+        q_size = num_head * head_dim
+        kv_size = num_head_kv * head_dim
+        q = torch.randn([batch_size, q_size], dtype=dtype)
+        k = torch.randn([batch_size, kv_size], dtype=dtype)
+        q_weight = torch.randn(head_dim, dtype=dtype)
+        k_weight = torch.randn(head_dim, dtype=dtype)
+        position_ids = torch.arange(batch_size, dtype=torch.int32)
+        base = 10000.0
+
+        ref_q = q.clone()
+        ref_k = k.clone()
+        ref_q = self._norm_per_head_native(ref_q, q_weight, head_dim, eps)
+        ref_k = self._norm_per_head_native(ref_k, k_weight, head_dim, eps)
+
+        def apply_rope(x: torch.Tensor, pos: int) -> torch.Tensor:
+            x = x.reshape(x.shape[0], -1, head_dim)
+            rotated = x.clone()
+            for b in range(x.shape[0]):
+                for h in range(x.shape[1]):
+                    row = x[b, h].clone()
+                    if is_neox:
+                        half = head_dim // 2
+                        for d in range(half):
+                            x0 = row[d]
+                            y0 = row[d + half]
+                            freq = base ** (-2.0 * d / head_dim)
+                            theta = pos * freq
+                            s = torch.sin(torch.tensor(theta, dtype=torch.float32))
+                            c = torch.cos(torch.tensor(theta, dtype=torch.float32))
+                            rotated[b, h, d] = x0 * c - y0 * s
+                            rotated[b, h, d + half] = y0 * c + x0 * s
+                    else:
+                        for d in range(0, head_dim, 2):
+                            x0 = row[d]
+                            y0 = row[d + 1]
+                            freq = base ** (-2.0 * (d / 2) / head_dim)
+                            theta = pos * freq
+                            s = torch.sin(torch.tensor(theta, dtype=torch.float32))
+                            c = torch.cos(torch.tensor(theta, dtype=torch.float32))
+                            rotated[b, h, d] = x0 * c - y0 * s
+                            rotated[b, h, d + 1] = y0 * c + x0 * s
+            return rotated.reshape_as(x)
+
+        for pos_idx, pos in enumerate(position_ids.tolist()):
+            ref_q[pos_idx] = apply_rope(ref_q[pos_idx : pos_idx + 1], pos).reshape(-1)
+            ref_k[pos_idx] = apply_rope(ref_k[pos_idx : pos_idx + 1], pos).reshape(-1)
+
+        half = head_dim // 2
+        freqs = base ** (-2.0 * torch.arange(half, dtype=torch.float32) / head_dim)
+        theta = position_ids.to(torch.float32)[:, None] * freqs[None, :]
+        cos_sin_cache = torch.cat([torch.cos(theta), torch.sin(theta)], dim=-1).to(
+            dtype
+        )
+
+        torch.ops.sgl_kernel.fused_qk_norm_rope_cpu(
+            q,
+            k,
+            q_weight,
+            k_weight,
+            eps,
+            is_neox,
+            position_ids,
+            cos_sin_cache,
+            head_dim,
+        )
+
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(q, ref_q, atol=atol, rtol=rtol)
+        torch.testing.assert_close(k, ref_k, atol=atol, rtol=rtol)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))