From e687b8d6af397de04e61229b63ca5cbbc757dd5f Mon Sep 17 00:00:00 2001 From: Cui Lily Date: Tue, 15 Sep 2026 14:17:57 +0800 Subject: [PATCH] [CPU] Implement fused QK Norm and RoPE kernels (#37748) Signed-off-by: Cui, Lily --- python/sglang/kernels/aot/csrc/cpu/norm.cpp | 277 ++++++++++++++++++ .../aot/csrc/cpu/torch_extension_cpu.cpp | 20 ++ python/sglang/srt/models/qwen3_moe.py | 85 ++++-- python/sglang/srt/models/utils.py | 33 ++- test/registered/cpu/test_norm.py | 127 ++++++++ 5 files changed, 517 insertions(+), 25 deletions(-) diff --git a/python/sglang/kernels/aot/csrc/cpu/norm.cpp b/python/sglang/kernels/aot/csrc/cpu/norm.cpp index 058441f04..5dff2cb09 100644 --- a/python/sglang/kernels/aot/csrc/cpu/norm.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/norm.cpp @@ -1052,3 +1052,280 @@ std::tuple fused_qk_gemma_rmsnorm_with_gate_ }); return std::make_tuple(q_out, k_out, gate_out); } + +namespace { + +template +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; + using fVec = at::vec::Vectorized; + 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(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(x_fvec0 * scale_fvec * w_fvec0, x_fvec1 * scale_fvec * w_fvec1).store(data + d); + } + for (; d < D; ++d) { + data[d] = static_cast(static_cast(data[d]) * scale * static_cast(weight[d])); + } +} + +template +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(data, is_q ? q_weight : k_weight, head_dim, eps); + } + }); +} + +template +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::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(xy0, xy1).store(data + d); + } + for (; d < rotary_dim; d += 2) { + const float x = static_cast(data[d]); + const float y = static_cast(data[d + 1]); + const float c = static_cast(cache[d / 2]); + const float s = static_cast(cache[half_rotary + d / 2]); + data[d] = static_cast(x * c - y * s); + data[d + 1] = static_cast(y * c + x * s); + } +} + +template +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::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(out0, out1).store(data + d); + convert_from_float_ext(out2, out3).store(data + half_rotary + d); + } + for (; d < half_rotary; ++d) { + const float x = static_cast(data[d]); + const float y = static_cast(data[d + half_rotary]); + const float c = static_cast(cache[d]); + const float s = static_cast(cache[half_rotary + d]); + data[d] = static_cast(x * c - y * s); + data[d + half_rotary] = static_cast(y * c + x * s); + } +} + +template +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(data, weight, head_dim, eps); + + if (is_neox) { + fused_qk_norm_rope_apply_neox(data, cache_row, rotary_dim); + } else { + fused_qk_norm_rope_apply_interleaved(data, cache_row, rotary_dim); + } +} + +template +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( + 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(q_weight, {head_dim}, st); + CHECK_INPUT_SHAPE_DTYPE(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( + q.data_ptr(), + k.data_ptr(), + q_weight.data_ptr(), + k_weight.data_ptr(), + num_tokens, + q.size(1) / head_dim, + k.size(1) / head_dim, + head_dim, + q.stride(0), + k.stride(0), + static_cast(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(q_weight, {head_dim}, st); + CHECK_INPUT_SHAPE_DTYPE(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 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(); + 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(); + } + fused_qk_norm_rope_kernel_impl( + q.data_ptr(), + k.data_ptr(), + q_weight.data_ptr(), + k_weight.data_ptr(), + num_tokens, + q.size(1) / head_dim, + k.size(1) / head_dim, + head_dim, + q.stride(0), + k.stride(0), + static_cast(eps), + is_neox, + pos_ptr, + cos_sin_cache.data_ptr(), + rotary_dim); + }); +} diff --git a/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp b/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp index 0100d216b..a184c5091 100644 --- a/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp @@ -58,6 +58,20 @@ at::Tensor fused_add_layernorm_cpu( const std::optional& 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 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)"); diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index a2a71be21..c0e93c371 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/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 diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 4c8ade007..e8a03a2f5 100644 --- a/python/sglang/srt/models/utils.py +++ b/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 diff --git a/test/registered/cpu/test_norm.py b/test/registered/cpu/test_norm.py index cd8c340ed..35652d0e6 100644 --- a/test/registered/cpu/test_norm.py +++ b/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__]))