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

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

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__]))