diff --git a/build.toml b/build.toml
index 3872776019c41eb700c91e3bf907e8b71bd47b25..a011dc8ffc58cb3726db1bafb34e31e473be8bb2 100644
--- a/build.toml
+++ b/build.toml
@@ -9,3 +9,11 @@ src = ["torch-ext/torch_binding.cpp"]
 backend = "cuda"
 depends = ["torch"]
 src = ["rotary/rotary_cuda.cu"]
+
+[kernel.rotary_xpu]
+backend = "xpu"
+depends = ["torch"]
+src = [
+    "rotary-xpu/rotary_xpu.cpp",
+    "rotary-xpu/rotary_xpu.hpp",
+    ]
\ No newline at end of file
diff --git a/build/torch27-cxx11-cu118-x86_64-linux/rotary/__init__.py b/build/torch27-cxx11-cu118-x86_64-linux/rotary/__init__.py
old mode 100644
new mode 100755
diff --git a/build/torch27-cxx11-cu118-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc b/build/torch27-cxx11-cu118-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..3e8cb222776a2da8e112221fdbb99a414d622cf4
Binary files /dev/null and b/build/torch27-cxx11-cu118-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc differ
diff --git a/build/torch27-cxx11-cu118-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc b/build/torch27-cxx11-cu118-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..05bd2d15b6eca322d474442e211e7e4d63071f50
Binary files /dev/null and b/build/torch27-cxx11-cu118-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc differ
diff --git a/build/torch27-cxx11-cu118-x86_64-linux/rotary/_ops.py b/build/torch27-cxx11-cu118-x86_64-linux/rotary/_ops.py
old mode 100644
new mode 100755
index 8d9717343c07cf81e45646b6fc80dddc95d58bdf..882a9998d9ad3e0226135704ffdc37207e41c499
--- a/build/torch27-cxx11-cu118-x86_64-linux/rotary/_ops.py
+++ b/build/torch27-cxx11-cu118-x86_64-linux/rotary/_ops.py
@@ -1,9 +1,9 @@
 import torch
-from . import _rotary_6b8e81d
-ops = torch.ops._rotary_6b8e81d
+from . import _rotary_cd1413b_dirty
+ops = torch.ops._rotary_cd1413b_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_rotary_6b8e81d::{op_name}"
\ No newline at end of file
+    return f"_rotary_cd1413b_dirty::{op_name}"
\ No newline at end of file
diff --git a/build/torch27-cxx11-cu118-x86_64-linux/rotary/_rotary_6b8e81d.abi3.so b/build/torch27-cxx11-cu118-x86_64-linux/rotary/_rotary_6b8e81d.abi3.so
deleted file mode 100755
index 94c6d44c7dba5b009caa43e9baffe81adedf970b..0000000000000000000000000000000000000000
--- a/build/torch27-cxx11-cu118-x86_64-linux/rotary/_rotary_6b8e81d.abi3.so
+++ /dev/null
@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e67a5779587b10616a0e19961ae50495c367f53da7c4a40a1c9b1f557537441d
-size 6287712
diff --git a/build/torch27-cxx11-cu118-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so b/build/torch27-cxx11-cu118-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
new file mode 100755
index 0000000000000000000000000000000000000000..2f86c8063c7fe14d1750fa98cf6180d4fafddd3e
--- /dev/null
+++ b/build/torch27-cxx11-cu118-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51c8d8635b97b599a33ba169458b47b9276f673c678c413107a5cab5a835f90e
+size 6807672
diff --git a/build/torch27-cxx11-cu126-x86_64-linux/rotary/__init__.py b/build/torch27-cxx11-cu126-x86_64-linux/rotary/__init__.py
old mode 100644
new mode 100755
diff --git a/build/torch27-cxx11-cu126-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc b/build/torch27-cxx11-cu126-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..dba38f5c4fda815a7a0c08a9be238e15b8b7dc82
Binary files /dev/null and b/build/torch27-cxx11-cu126-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc differ
diff --git a/build/torch27-cxx11-cu126-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc b/build/torch27-cxx11-cu126-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..eca86af4da07beef6c418ed60c34ba89b676ab59
Binary files /dev/null and b/build/torch27-cxx11-cu126-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc differ
diff --git a/build/torch27-cxx11-cu126-x86_64-linux/rotary/_ops.py b/build/torch27-cxx11-cu126-x86_64-linux/rotary/_ops.py
old mode 100644
new mode 100755
index 8d9717343c07cf81e45646b6fc80dddc95d58bdf..882a9998d9ad3e0226135704ffdc37207e41c499
--- a/build/torch27-cxx11-cu126-x86_64-linux/rotary/_ops.py
+++ b/build/torch27-cxx11-cu126-x86_64-linux/rotary/_ops.py
@@ -1,9 +1,9 @@
 import torch
-from . import _rotary_6b8e81d
-ops = torch.ops._rotary_6b8e81d
+from . import _rotary_cd1413b_dirty
+ops = torch.ops._rotary_cd1413b_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_rotary_6b8e81d::{op_name}"
\ No newline at end of file
+    return f"_rotary_cd1413b_dirty::{op_name}"
\ No newline at end of file
diff --git a/build/torch27-cxx11-cu126-x86_64-linux/rotary/_rotary_6b8e81d.abi3.so b/build/torch27-cxx11-cu126-x86_64-linux/rotary/_rotary_6b8e81d.abi3.so
deleted file mode 100755
index 66c3721a4eb80c7d99bb53766a1c66cb9b4303f6..0000000000000000000000000000000000000000
--- a/build/torch27-cxx11-cu126-x86_64-linux/rotary/_rotary_6b8e81d.abi3.so
+++ /dev/null
@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ced538b9a02abba45b7769dd09902cbc1816091ec50da2b65c0549f85974ea4c
-size 6296464
diff --git a/build/torch27-cxx11-cu126-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so b/build/torch27-cxx11-cu126-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
new file mode 100755
index 0000000000000000000000000000000000000000..43f652168a6298256b697f26a9bd1e53108dd855
--- /dev/null
+++ b/build/torch27-cxx11-cu126-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b49a2fb4c22c6cda6d4d28d1f5eb3ad84801174c7790628519f0c7529a57773
+size 6820520
diff --git a/build/torch27-cxx11-cu128-x86_64-linux/rotary/__init__.py b/build/torch27-cxx11-cu128-x86_64-linux/rotary/__init__.py
old mode 100644
new mode 100755
diff --git a/build/torch27-cxx11-cu128-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc b/build/torch27-cxx11-cu128-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..a6e3f79acb019fbf14ac6ffd999448170ef1b4d5
Binary files /dev/null and b/build/torch27-cxx11-cu128-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc differ
diff --git a/build/torch27-cxx11-cu128-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc b/build/torch27-cxx11-cu128-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..ba5d3a9275df973397473625f68d37ff53f76eea
Binary files /dev/null and b/build/torch27-cxx11-cu128-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc differ
diff --git a/build/torch27-cxx11-cu128-x86_64-linux/rotary/_ops.py b/build/torch27-cxx11-cu128-x86_64-linux/rotary/_ops.py
old mode 100644
new mode 100755
index 8d9717343c07cf81e45646b6fc80dddc95d58bdf..882a9998d9ad3e0226135704ffdc37207e41c499
--- a/build/torch27-cxx11-cu128-x86_64-linux/rotary/_ops.py
+++ b/build/torch27-cxx11-cu128-x86_64-linux/rotary/_ops.py
@@ -1,9 +1,9 @@
 import torch
-from . import _rotary_6b8e81d
-ops = torch.ops._rotary_6b8e81d
+from . import _rotary_cd1413b_dirty
+ops = torch.ops._rotary_cd1413b_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_rotary_6b8e81d::{op_name}"
\ No newline at end of file
+    return f"_rotary_cd1413b_dirty::{op_name}"
\ No newline at end of file
diff --git a/build/torch27-cxx11-cu128-x86_64-linux/rotary/_rotary_6b8e81d.abi3.so b/build/torch27-cxx11-cu128-x86_64-linux/rotary/_rotary_6b8e81d.abi3.so
deleted file mode 100755
index e3be43836a3f46b11493b2a64eeb33e560b7bd9c..0000000000000000000000000000000000000000
--- a/build/torch27-cxx11-cu128-x86_64-linux/rotary/_rotary_6b8e81d.abi3.so
+++ /dev/null
@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e4a64022928f62fe0b476c54f5a426b856296fc2bba796c7e9fe6406c6d65485
-size 10157336
diff --git a/build/torch27-cxx11-cu128-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so b/build/torch27-cxx11-cu128-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
new file mode 100755
index 0000000000000000000000000000000000000000..2a50c157d413563fcba486b8104bea699d4f05e1
--- /dev/null
+++ b/build/torch27-cxx11-cu128-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:638bbc069d927f9e37f1e720e73ee4af097ec16fc882f7abcc04dae2045b80a1
+size 10529832
diff --git a/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/__init__.py b/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..eba8039e210c8b710c5c663ef4e7930757f271be
--- /dev/null
+++ b/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/__init__.py
@@ -0,0 +1,19 @@
+from typing import Tuple
+import torch
+
+from ._ops import ops
+
+
+def apply_rotary(
+    x1: torch.Tensor,
+    x2: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out1: torch.Tensor,
+    out2: torch.Tensor,
+    conj: bool,
+):
+    ops.apply_rotary(x1, x2, cos, sin, out1, out2, conj)
+
+
+__all__ = ["apply_rotary"]
diff --git a/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc b/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..abd6c81dc9fecc7d53eb660acfea14d138cb3936
Binary files /dev/null and b/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc differ
diff --git a/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc b/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..5ffe3824a7c006aaad48f0acad988ac0ffdd0359
Binary files /dev/null and b/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc differ
diff --git a/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/_ops.py b/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/_ops.py
new file mode 100755
index 0000000000000000000000000000000000000000..882a9998d9ad3e0226135704ffdc37207e41c499
--- /dev/null
+++ b/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/_ops.py
@@ -0,0 +1,9 @@
+import torch
+from . import _rotary_cd1413b_dirty
+ops = torch.ops._rotary_cd1413b_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_rotary_cd1413b_dirty::{op_name}"
\ No newline at end of file
diff --git a/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so b/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
new file mode 100755
index 0000000000000000000000000000000000000000..a30dcc3d3d7f34c0bd8a5500a0c5b88d700664ff
--- /dev/null
+++ b/build/torch27-cxx11-xpu20250-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7f449e098ab5bbe9ca35e2a904132ed4e378e54d579aefe95e4e83e07a73bfe
+size 2248696
diff --git a/build/torch28-cxx11-cu126-x86_64-linux/rotary/__init__.py b/build/torch28-cxx11-cu126-x86_64-linux/rotary/__init__.py
old mode 100644
new mode 100755
diff --git a/build/torch28-cxx11-cu126-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc b/build/torch28-cxx11-cu126-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..51fe290a3984f6fb962af5977b274a4f09bc5979
Binary files /dev/null and b/build/torch28-cxx11-cu126-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc differ
diff --git a/build/torch28-cxx11-cu126-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc b/build/torch28-cxx11-cu126-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..b905458ddf063c5ee2a88c9d665ea0679237836a
Binary files /dev/null and b/build/torch28-cxx11-cu126-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc differ
diff --git a/build/torch28-cxx11-cu126-x86_64-linux/rotary/_ops.py b/build/torch28-cxx11-cu126-x86_64-linux/rotary/_ops.py
old mode 100644
new mode 100755
index dd4baa9daad5b83b5309896a2982c812d8c389b1..882a9998d9ad3e0226135704ffdc37207e41c499
--- a/build/torch28-cxx11-cu126-x86_64-linux/rotary/_ops.py
+++ b/build/torch28-cxx11-cu126-x86_64-linux/rotary/_ops.py
@@ -1,9 +1,9 @@
 import torch
-from . import _rotary_d5e8892
-ops = torch.ops._rotary_d5e8892
+from . import _rotary_cd1413b_dirty
+ops = torch.ops._rotary_cd1413b_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_rotary_d5e8892::{op_name}"
\ No newline at end of file
+    return f"_rotary_cd1413b_dirty::{op_name}"
\ No newline at end of file
diff --git a/build/torch28-cxx11-cu126-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so b/build/torch28-cxx11-cu126-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
new file mode 100755
index 0000000000000000000000000000000000000000..71b6eaf4f9e54c8f1a01b2e2d01dc88382a8650e
--- /dev/null
+++ b/build/torch28-cxx11-cu126-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ea610dc89fe7d0037e55ff171158e099f79a08fcd14a8aff117c7b090d79a6e2
+size 6817216
diff --git a/build/torch28-cxx11-cu126-x86_64-linux/rotary/_rotary_d5e8892.abi3.so b/build/torch28-cxx11-cu126-x86_64-linux/rotary/_rotary_d5e8892.abi3.so
deleted file mode 100755
index 13479c760691c87133542d81e6dccf2a7cbf6bf8..0000000000000000000000000000000000000000
--- a/build/torch28-cxx11-cu126-x86_64-linux/rotary/_rotary_d5e8892.abi3.so
+++ /dev/null
@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:663ac605052413c175dcff4dec8545aa37ad67e8629002621741b89b182d98cd
-size 6293168
diff --git a/build/torch28-cxx11-cu128-x86_64-linux/rotary/__init__.py b/build/torch28-cxx11-cu128-x86_64-linux/rotary/__init__.py
old mode 100644
new mode 100755
diff --git a/build/torch28-cxx11-cu128-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc b/build/torch28-cxx11-cu128-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..6710462bf948980a76edb3587ca3260eb7e93cb8
Binary files /dev/null and b/build/torch28-cxx11-cu128-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc differ
diff --git a/build/torch28-cxx11-cu128-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc b/build/torch28-cxx11-cu128-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..932be89ecdd58624ee42ade87d921fe9beeacef0
Binary files /dev/null and b/build/torch28-cxx11-cu128-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc differ
diff --git a/build/torch28-cxx11-cu128-x86_64-linux/rotary/_ops.py b/build/torch28-cxx11-cu128-x86_64-linux/rotary/_ops.py
old mode 100644
new mode 100755
index dd4baa9daad5b83b5309896a2982c812d8c389b1..882a9998d9ad3e0226135704ffdc37207e41c499
--- a/build/torch28-cxx11-cu128-x86_64-linux/rotary/_ops.py
+++ b/build/torch28-cxx11-cu128-x86_64-linux/rotary/_ops.py
@@ -1,9 +1,9 @@
 import torch
-from . import _rotary_d5e8892
-ops = torch.ops._rotary_d5e8892
+from . import _rotary_cd1413b_dirty
+ops = torch.ops._rotary_cd1413b_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_rotary_d5e8892::{op_name}"
\ No newline at end of file
+    return f"_rotary_cd1413b_dirty::{op_name}"
\ No newline at end of file
diff --git a/build/torch28-cxx11-cu128-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so b/build/torch28-cxx11-cu128-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
new file mode 100755
index 0000000000000000000000000000000000000000..68f111fb1871eca79bbbb017cc4b6ea94c7e6165
--- /dev/null
+++ b/build/torch28-cxx11-cu128-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f8af2b94a8121a7a8b6ac446cb6eb117d49cb4ea8842d7024bb1b9b26fb97db
+size 10526424
diff --git a/build/torch28-cxx11-cu128-x86_64-linux/rotary/_rotary_d5e8892.abi3.so b/build/torch28-cxx11-cu128-x86_64-linux/rotary/_rotary_d5e8892.abi3.so
deleted file mode 100755
index f3ada6a304ed16bf9a2af9af55ecdaa3b4b2eb95..0000000000000000000000000000000000000000
--- a/build/torch28-cxx11-cu128-x86_64-linux/rotary/_rotary_d5e8892.abi3.so
+++ /dev/null
@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:264410fa13bca33e706d1c3eb12a2d966e8fa07e2b786cbd8332d462f4883d1a
-size 10149824
diff --git a/build/torch28-cxx11-cu129-x86_64-linux/rotary/__init__.py b/build/torch28-cxx11-cu129-x86_64-linux/rotary/__init__.py
old mode 100644
new mode 100755
diff --git a/build/torch28-cxx11-cu129-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc b/build/torch28-cxx11-cu129-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..3c0ebd8f6fd6ceaa1d5ebc8d7c8e875cda718787
Binary files /dev/null and b/build/torch28-cxx11-cu129-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc differ
diff --git a/build/torch28-cxx11-cu129-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc b/build/torch28-cxx11-cu129-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..6ced4aee44fb96b785ecdb9bbfed7dfdebc7497d
Binary files /dev/null and b/build/torch28-cxx11-cu129-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc differ
diff --git a/build/torch28-cxx11-cu129-x86_64-linux/rotary/_ops.py b/build/torch28-cxx11-cu129-x86_64-linux/rotary/_ops.py
old mode 100644
new mode 100755
index dd4baa9daad5b83b5309896a2982c812d8c389b1..882a9998d9ad3e0226135704ffdc37207e41c499
--- a/build/torch28-cxx11-cu129-x86_64-linux/rotary/_ops.py
+++ b/build/torch28-cxx11-cu129-x86_64-linux/rotary/_ops.py
@@ -1,9 +1,9 @@
 import torch
-from . import _rotary_d5e8892
-ops = torch.ops._rotary_d5e8892
+from . import _rotary_cd1413b_dirty
+ops = torch.ops._rotary_cd1413b_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_rotary_d5e8892::{op_name}"
\ No newline at end of file
+    return f"_rotary_cd1413b_dirty::{op_name}"
\ No newline at end of file
diff --git a/build/torch28-cxx11-cu129-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so b/build/torch28-cxx11-cu129-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
new file mode 100755
index 0000000000000000000000000000000000000000..17440b83122c4aa653435d73d771d12eac66b996
--- /dev/null
+++ b/build/torch28-cxx11-cu129-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88098a942da850ef34bc5d4b2f810d9c3092718c134fba911161a04eba73c559
+size 10586840
diff --git a/build/torch28-cxx11-cu129-x86_64-linux/rotary/_rotary_d5e8892.abi3.so b/build/torch28-cxx11-cu129-x86_64-linux/rotary/_rotary_d5e8892.abi3.so
deleted file mode 100755
index eec1aac04e680c4e7bbf7f8fd19ff12436015a81..0000000000000000000000000000000000000000
--- a/build/torch28-cxx11-cu129-x86_64-linux/rotary/_rotary_d5e8892.abi3.so
+++ /dev/null
@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:dd62535c2713d05e74f4b53c84564caeeba51aaf06f5fe59a3182b04a5ae3c5a
-size 10169280
diff --git a/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__init__.py b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..eba8039e210c8b710c5c663ef4e7930757f271be
--- /dev/null
+++ b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__init__.py
@@ -0,0 +1,19 @@
+from typing import Tuple
+import torch
+
+from ._ops import ops
+
+
+def apply_rotary(
+    x1: torch.Tensor,
+    x2: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out1: torch.Tensor,
+    out2: torch.Tensor,
+    conj: bool,
+):
+    ops.apply_rotary(x1, x2, cos, sin, out1, out2, conj)
+
+
+__all__ = ["apply_rotary"]
diff --git a/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__pycache__/__init__.cpython-311.pyc b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..54f458b69bed5dbd698dfa26fe00c0b02c914e4a
Binary files /dev/null and b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__pycache__/__init__.cpython-311.pyc differ
diff --git a/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..c3534f51269371cdca09df12b23130847cd60863
Binary files /dev/null and b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__pycache__/__init__.cpython-313.pyc differ
diff --git a/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__pycache__/_ops.cpython-311.pyc b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__pycache__/_ops.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c85210970cad2446f178ffdebe2715fb1925fd97
Binary files /dev/null and b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__pycache__/_ops.cpython-311.pyc differ
diff --git a/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc
new file mode 100755
index 0000000000000000000000000000000000000000..12f1525244249e7c96804944827577119a4c6e0c
Binary files /dev/null and b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/__pycache__/_ops.cpython-313.pyc differ
diff --git a/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/_ops.py b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/_ops.py
new file mode 100755
index 0000000000000000000000000000000000000000..882a9998d9ad3e0226135704ffdc37207e41c499
--- /dev/null
+++ b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/_ops.py
@@ -0,0 +1,9 @@
+import torch
+from . import _rotary_cd1413b_dirty
+ops = torch.ops._rotary_cd1413b_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_rotary_cd1413b_dirty::{op_name}"
\ No newline at end of file
diff --git a/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
new file mode 100755
index 0000000000000000000000000000000000000000..6f020bf2e546d40af6f3d01c393c7911a2df7642
--- /dev/null
+++ b/build/torch28-cxx11-xpu20251-x86_64-linux/rotary/_rotary_cd1413b_dirty.abi3.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:12dba600201d5c2bd5cd123afc3b65f835f3698a89b75e61c85ee3f359f2e901
+size 2239816
diff --git a/flake.lock b/flake.lock
index 26319b67eda604001ec12fc6bdd52ee9bfe65846..9a10662457740e93808381972c1b39c46faea045 100644
--- a/flake.lock
+++ b/flake.lock
@@ -17,11 +17,11 @@
     },
     "flake-compat_2": {
       "locked": {
-        "lastModified": 1733328505,
-        "narHash": "sha256-NeCCThCEP3eCl2l/+27kNNK7QrwZB1IJCrXfrbv5oqU=",
+        "lastModified": 1747046372,
+        "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=",
         "owner": "edolstra",
         "repo": "flake-compat",
-        "rev": "ff81ac966bb2cae68946d5ed5fc4994f96d0ffec",
+        "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885",
         "type": "github"
       },
       "original": {
@@ -73,11 +73,11 @@
         "nixpkgs": "nixpkgs"
       },
       "locked": {
-        "lastModified": 1753354560,
-        "narHash": "sha256-vmOfRmr0Qm/IbZTWB2sBn+UFrABSTTA/cTg+m27Yt/E=",
+        "lastModified": 1757493151,
+        "narHash": "sha256-eirWlcvs2rjZmU8JcF4CKN1IEnNfpQnGuf2qbK3IQh8=",
         "owner": "huggingface",
         "repo": "hf-nix",
-        "rev": "7f2aceda2a2e72cd573bdb25e5c0667fd75f89d3",
+        "rev": "503cd4eb9866103c983dbef93d9ad5db4fb6b415",
         "type": "github"
       },
       "original": {
@@ -98,33 +98,32 @@
         ]
       },
       "locked": {
-        "lastModified": 1753354632,
-        "narHash": "sha256-31SX3Raiyx0qCuY9JSlx9ZZgxljeUxvW+JdujjxbofQ=",
+        "lastModified": 1757570810,
+        "narHash": "sha256-YFWQwy2LKbhjdLW8wkyNkE/+Vbdn6qlJif2CKvBT9Qo=",
         "owner": "huggingface",
         "repo": "kernel-builder",
-        "rev": "524b628fd8e58525dbd28455bffb0628092c5265",
+        "rev": "1201847af3ff757b65015c6e06b5bd75896d2d4b",
         "type": "github"
       },
       "original": {
         "owner": "huggingface",
-        "ref": "torch-2.8",
         "repo": "kernel-builder",
         "type": "github"
       }
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1752785354,
-        "narHash": "sha256-Y33ryUz7MPqKrZwlbQcsYCUz2jAJCacRf8jbs0tYUlA=",
+        "lastModified": 1755963616,
+        "narHash": "sha256-6yD0ww/S8n+U2uPYcJZ3DRURP8Kx036GRpR2uPNZroE=",
         "owner": "nixos",
         "repo": "nixpkgs",
-        "rev": "d38025438a6ee456758dc03188ca6873a415463b",
+        "rev": "73e96df7cff5783f45e21342a75a1540c4eddce4",
         "type": "github"
       },
       "original": {
         "owner": "nixos",
+        "ref": "nixos-unstable-small",
         "repo": "nixpkgs",
-        "rev": "d38025438a6ee456758dc03188ca6873a415463b",
         "type": "github"
       }
     },
diff --git a/flake.nix b/flake.nix
index f655a849d1d070474437c8ee846db3dc4b0b6965..a4c72b81489bcc67dff0d5a2857d765a0a74c5b0 100644
--- a/flake.nix
+++ b/flake.nix
@@ -1,15 +1,9 @@
 {
-  description = "Flake for rotary kernel";
-
+  description = "Flake for Torch kernel extension";
   inputs = {
-    kernel-builder.url = "github:huggingface/kernel-builder/torch-2.8";
+    kernel-builder.url = "github:huggingface/kernel-builder";
   };
-
-  outputs =
-    {
-      self,
-      kernel-builder,
-    }:
+  outputs = { self, kernel-builder, }:
     kernel-builder.lib.genFlakeOutputs {
       path = ./.;
       rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate;
diff --git a/rotary-xpu/rotary_xpu.cpp b/rotary-xpu/rotary_xpu.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..4d30083aeef5360cc129fd4abaacdb902d91dca9
--- /dev/null
+++ b/rotary-xpu/rotary_xpu.cpp
@@ -0,0 +1,40 @@
+#include <torch/all.h>
+#include "rotary_xpu.hpp"
+
+void _apply_rotary(torch::Tensor const &x1, torch::Tensor const &x2,
+                       torch::Tensor const &cos, torch::Tensor const &sin,
+                       torch::Tensor &out1, torch::Tensor &out2,
+                       bool const conj) {
+    auto iter = at::TensorIteratorConfig()
+        .add_output(out1)
+        .add_output(out2)
+        .add_input(x1)
+        .add_input(x2)
+        .add_input(cos)
+        .add_input(sin)
+        .check_all_same_dtype(false)
+        .promote_inputs_to_common_dtype(false)
+        .build();
+
+    if (!conj) {
+        AT_DISPATCH_FLOATING_TYPES_AND2(at::kBFloat16, at::kHalf, x1.scalar_type(), "rotary_kernel_xpu", [&] {
+            gpu_kernel_multiple_outputs(
+                iter, [] (scalar_t x1, scalar_t x2, scalar_t cos,
+                                    scalar_t sin) -> std::tuple<scalar_t, scalar_t> {
+                scalar_t out1 = float(x1) * float(cos) - float(x2) * float(sin);
+                scalar_t out2 = float(x1) * float(sin) + float(x2) * float(cos);
+                return {out1, out2};
+            });
+        });
+    } else {
+        AT_DISPATCH_FLOATING_TYPES_AND2(at::kBFloat16, at::kHalf, x1.scalar_type(), "rotary_kernel_xpu", [&] {
+            gpu_kernel_multiple_outputs(
+                iter, [] (scalar_t x1, scalar_t x2, scalar_t cos,
+                                    scalar_t sin) -> std::tuple<scalar_t, scalar_t> {
+                scalar_t out1 = float(x1) * float(cos) + float(x2) * float(sin);
+                scalar_t out2 = -float(x1) * float(sin) + float(x2) * float(cos);
+                return {out1, out2};
+            });
+        });
+    }
+}
diff --git a/rotary-xpu/rotary_xpu.hpp b/rotary-xpu/rotary_xpu.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..5f0f6e24055695c710c86fbc6208907d160f36e2
--- /dev/null
+++ b/rotary-xpu/rotary_xpu.hpp
@@ -0,0 +1,375 @@
+#include <ATen/core/TensorBody.h>
+#include <ATen/detail/FunctionTraits.h>
+#include <ATen/native/TensorIterator.h>
+#include <sycl/sycl.hpp>
+#include <ATen/core/Array.h>
+#include <c10/macros/Macros.h>
+#include <c10/util/Exception.h>
+#include <c10/util/TypeCast.h>
+#include <cstdint>
+#include <type_traits>
+#include <array>
+#include <c10/core/ScalarType.h>
+#include <c10/xpu/XPUStream.h>
+#include <ATen/xpu/XPUContext.h>
+
+constexpr int MAX_DIMS = 12;
+
+struct LoadWithoutCast {
+  template <typename scalar_t>
+  C10_DEVICE scalar_t load(char* base_ptr, uint32_t offset, int arg) {
+    return c10::load(reinterpret_cast<scalar_t*>(base_ptr) + offset);
+  }
+};
+
+struct StoreWithoutCast {
+  template <typename scalar_t>
+  C10_DEVICE void store(scalar_t value, char* base_ptr, uint32_t offset, int arg = 0) {
+    *(reinterpret_cast<scalar_t*>(base_ptr) + offset) = value;
+  }
+};
+
+template <template <int i> typename func, int end, int current = 0>
+struct static_unroll {
+  template <typename... Args>
+  static inline C10_HOST_DEVICE void with_args(Args&&... args) {
+    func<current>::apply(std::forward<Args>(args)...);
+    static_unroll<func, end, current + 1>::with_args(args...);
+  }
+};
+
+template <template <int i> typename func, int end>
+struct static_unroll<func, end, end> {
+  template <typename... Args>
+  static inline C10_HOST_DEVICE void with_args(Args... args) {}
+};
+
+template <int current>
+struct multi_outputs_store_helper {
+  template <int ntensors, int num_outputs, typename... Args>
+  static C10_HOST_DEVICE void apply(
+      at::detail::Array<char*, ntensors> data,
+      at::detail::Array<uint32_t, num_outputs> offsets,
+      std::tuple<Args...> ret) {
+    using T = typename std::tuple_element<current, std::tuple<Args...>>::type;
+    T* to = reinterpret_cast<T*>(data[current]) + offsets[current];
+    *to = std::get<current>(ret);
+  }
+};
+
+template <int arg_index>
+struct unroll_load_helper {
+  template <typename args_t, typename policy_t, typename offset_t, typename loader_t>
+  static C10_DEVICE void apply(
+      policy_t& self,
+      args_t* args,
+      offset_t offset,
+      loader_t loader,
+      int j,
+      int num_outputs) {
+    using arg_t = std::tuple_element_t<arg_index, args_t>;
+    std::get<arg_index>(args[j]) = loader.template load<arg_t>(
+        self.data[arg_index + num_outputs], offset[arg_index], arg_index);
+  }
+};
+
+template <int item_work_size, typename data_t, typename inp_calc_t, typename out_calc_t, int num_outputs>
+struct multi_outputs_unroll {
+  data_t data;
+  int remaining;
+  inp_calc_t input_offset_calculator;
+  out_calc_t output_offset_calculator;
+  LoadWithoutCast loader;
+  StoreWithoutCast storer;
+  int item_idx;
+  int group_idx;
+  int num_items_per_group;
+  int group_work_size;
+
+  multi_outputs_unroll(
+      data_t data,
+      int remaining,
+      inp_calc_t ic,
+      out_calc_t oc,
+      int item_idx,
+      int group_idx,
+      int num_items_per_group)
+      : data(data),
+        remaining(remaining),
+        input_offset_calculator(ic),
+        output_offset_calculator(oc),
+        item_idx(item_idx),
+        group_idx(group_idx),
+        num_items_per_group(num_items_per_group),
+        group_work_size(item_work_size * num_items_per_group) {}
+
+  inline bool check_inbounds(int item_work_elem) const {
+    return (item_idx + item_work_elem * num_items_per_group < remaining);
+  }
+
+  template <typename args_t>
+  inline void load(args_t* args) {
+    constexpr int arity = std::tuple_size<args_t>::value;
+    int item_idx_ = item_idx;
+#pragma unroll
+    for (int i = 0; i < item_work_size; i++) {
+      if (item_idx_ >= remaining) {
+        return;
+      }
+      int linear_idx = item_idx_ + group_work_size * group_idx;
+      auto offset = input_offset_calculator.get(linear_idx);
+      static_unroll<unroll_load_helper, arity>::with_args(
+          *this, args, offset, loader, i, num_outputs);
+      item_idx_ += num_items_per_group;
+    }
+  }
+
+  template <typename return_t>
+  inline void store(return_t* from) {
+    int item_idx_ = item_idx;
+#pragma unroll
+    for (int i = 0; i < item_work_size; i++) {
+      if (item_idx_ >= this->remaining) {
+        return;
+      }
+      int linear_idx = item_idx_ + group_work_size * group_idx;
+      auto offsets = this->output_offset_calculator.get(linear_idx);
+      static_unroll<multi_outputs_store_helper, num_outputs>::with_args(this->data, offsets, from[i]);
+      item_idx_ += num_items_per_group;
+    }
+  }
+};
+
+template <int item_work_size, typename func_t, typename policy_t>
+inline void elementwise_kernel_helper(func_t f, policy_t policy) {
+  using traits = function_traits<func_t>;
+  using return_t = typename traits::result_type;
+  using args_t = typename traits::ArgsTuple;
+
+  return_t results[item_work_size];
+  args_t args[item_work_size];
+
+  policy.load(args);
+
+#pragma unroll
+  for (int i = 0; i < item_work_size; i++) {
+    if (policy.check_inbounds(i)) {
+      results[i] = std::apply(f, args[i]);
+    }
+  }
+
+  policy.store(results);
+}
+
+template <int num_outputs, typename func_t, typename array_t, typename in_calc_t, typename out_calc_t>
+struct UnrolledElementwiseForMultiOutputsKernel {
+  static constexpr int item_work_size = 4;
+
+  void operator()(sycl::nd_item<1> item_id) const {
+    int grpsz = item_id.get_local_range(0);
+    int grpid = item_id.get_group(0);
+    int lid = item_id.get_local_id(0);
+    int remaining = numel_ - item_work_size * grpsz * grpid;
+    auto policy = multi_outputs_unroll<item_work_size, array_t, in_calc_t, out_calc_t, num_outputs>(
+        data_, remaining, ic_, oc_, lid, grpid, grpsz);
+    elementwise_kernel_helper<item_work_size>(f_, policy);
+  };
+
+  UnrolledElementwiseForMultiOutputsKernel(int numel, func_t f, array_t data, in_calc_t ic, out_calc_t oc)
+      : numel_(numel), f_(f), data_(data), ic_(ic), oc_(oc) {}
+
+ private:
+  int numel_;
+  func_t f_;
+  array_t data_;
+  in_calc_t ic_;
+  out_calc_t oc_;
+};
+
+template <typename Value>
+struct IntDivider {
+  IntDivider() = default;
+  IntDivider(Value d) : divisor(d) {}
+
+  C10_HOST_DEVICE inline Value div(Value n) const {
+    return n / divisor;
+  }
+  C10_HOST_DEVICE inline Value mod(Value n) const {
+    return n % divisor;
+  }
+  C10_HOST_DEVICE inline auto divmod(Value n) const {
+    return std::make_pair(n / divisor, n % divisor);
+  }
+
+  Value divisor;
+};
+
+template <int NARGS, typename index_t = uint32_t, bool signed_strides = false>
+struct OffsetCalculator {
+  using stride_t = std::conditional_t<signed_strides, std::make_signed_t<index_t>, index_t>;
+  using offset_type = at::detail::Array<stride_t, std::max<int>(NARGS, 1)>;
+
+  OffsetCalculator(int dims, const int64_t* sizes, const int64_t* const* strides, const int64_t* element_sizes = nullptr)
+      : dims(dims) {
+    TORCH_CHECK(dims <= MAX_DIMS, "tensor has too many (>", MAX_DIMS, ") dims");
+    for (int i = 0; i < dims; i++) {
+      sizes_[i] = IntDivider<index_t>(sizes[i]);
+      for (int arg = 0; arg < NARGS; arg++) {
+        int64_t element_size = (element_sizes == nullptr ? 1LL : element_sizes[arg]);
+        strides_[i][arg] = strides[arg][i] / element_size;
+      }
+    }
+  }
+
+  C10_HOST_DEVICE offset_type get(index_t linear_idx) const {
+    offset_type offsets;
+#pragma unroll
+    for (int arg = 0; arg < NARGS; arg++) {
+      offsets[arg] = 0;
+    }
+
+#pragma unroll
+    for (int dim = 0; dim < MAX_DIMS; ++dim) {
+      if (dim == dims) {
+        break;
+      }
+      auto divmod = sizes_[dim].divmod(linear_idx);
+      linear_idx = divmod.first;
+
+#pragma unroll
+      for (int arg = 0; arg < NARGS; arg++) {
+        offsets[arg] += divmod.second * strides_[dim][arg];
+      }
+    }
+    return offsets;
+  }
+
+  int dims;
+  IntDivider<index_t> sizes_[MAX_DIMS];
+  stride_t strides_[MAX_DIMS][std::max<int>(NARGS, 1)];
+};
+
+template <int N>
+static OffsetCalculator<N> make_input_offset_calculator(const at::TensorIteratorBase& iter) {
+  constexpr int array_size = std::max<int>(N, 1);
+  TORCH_INTERNAL_ASSERT(N == iter.ntensors() - iter.noutputs());
+  std::array<const int64_t*, array_size> strides;
+  int64_t element_sizes[array_size];
+  for (int i = 0; i < N; i++) {
+    strides[i] = iter.strides(i + iter.noutputs()).data();
+    element_sizes[i] = iter.element_size(i + iter.noutputs());
+  }
+  return OffsetCalculator<N>(iter.ndim(), iter.shape().data(), strides.data(), element_sizes);
+}
+
+template <int num_outputs = 1>
+static OffsetCalculator<num_outputs> make_output_offset_calculator(const at::TensorIteratorBase& iter) {
+  TORCH_INTERNAL_ASSERT(num_outputs == iter.noutputs());
+  std::array<const int64_t*, num_outputs> strides;
+  int64_t element_sizes[num_outputs];
+  for (int i = 0; i < num_outputs; i++) {
+    strides[i] = iter.strides(i).data();
+    element_sizes[i] = iter.element_size(i);
+  }
+  return OffsetCalculator<num_outputs>(iter.ndim(), iter.shape().data(), strides.data(), element_sizes);
+}
+
+static inline int64_t syclMaxWorkItemsPerSubSlice(at::DeviceIndex dev_id = c10::xpu::getCurrentXPUStream().device_index()) {
+  auto* dev_prop = at::xpu::getDeviceProperties(dev_id);
+  int64_t simd_width = dev_prop->sub_group_sizes[0];
+  int64_t eu_count = dev_prop->gpu_eu_count_per_subslice;
+  return simd_width * eu_count;
+}
+
+template<class T>
+T ceil_div(T dividend, T divisor) {
+    return (dividend + divisor - 1) / divisor;
+}
+
+template <typename ker_t>
+static inline void sycl_kernel_submit(int64_t global_range, int64_t local_range, ::sycl::queue q, ker_t ker) {
+  q.parallel_for(
+    sycl::nd_range<1>(sycl::range<1>(global_range), sycl::range<1>(local_range)),
+    ker
+  );
+}
+
+template <int num_outputs, typename func_t, typename array_t, typename in_calc_t, typename out_calc_t>
+static inline void launch_unrolled_kernel_for_multi_outputs(
+    int64_t N,
+    const func_t& f,
+    array_t data,
+    in_calc_t ic,
+    out_calc_t oc) {
+  TORCH_INTERNAL_ASSERT(N > 0 && N <= std::numeric_limits<int32_t>::max());
+
+  auto ker = UnrolledElementwiseForMultiOutputsKernel<num_outputs, func_t, array_t, in_calc_t, out_calc_t>(N, f, data, ic, oc);
+  using ker_t = decltype(ker);
+
+  int wg_sz = syclMaxWorkItemsPerSubSlice();
+  int num_wg = ceil_div<int>(N, ker_t::item_work_size * wg_sz);
+  sycl_kernel_submit(wg_sz * num_wg, wg_sz, c10::xpu::getCurrentXPUStream().queue(), ker);
+}
+
+template <int N>
+struct TrivialOffsetCalculator {
+  using offset_type = at::detail::Array<uint32_t, std::max<int>(N, 1)>;
+
+  C10_HOST_DEVICE offset_type get(uint32_t linear_idx) const {
+    offset_type offsets;
+#pragma unroll
+    for (int arg = 0; arg < N; arg++) {
+      offsets[arg] = linear_idx;
+    }
+    return offsets;
+  }
+};
+
+template <typename func_t>
+void gpu_kernel_multiple_outputs_impl(at::TensorIteratorBase& iter, const func_t& f) {
+  using traits = function_traits<func_t>;
+  using output_t = typename traits::result_type;
+  constexpr int num_outputs = std::tuple_size<output_t>::value;
+  constexpr int num_inputs = traits::arity;
+  constexpr int ntensors = num_outputs + num_inputs;
+
+  TORCH_INTERNAL_ASSERT(iter.can_use_32bit_indexing());
+  TORCH_INTERNAL_ASSERT(iter.ntensors() == ntensors);
+
+  at::detail::Array<char*, ntensors> data;
+  for (int i = 0; i < ntensors; i++) {
+    data[i] = (char*)iter.data_ptr(i);
+  }
+
+  int64_t numel = iter.numel();
+
+  if (iter.is_contiguous()) {
+    auto input_calc = TrivialOffsetCalculator<num_inputs>();
+    auto output_calc = TrivialOffsetCalculator<num_outputs>();
+    launch_unrolled_kernel_for_multi_outputs<num_outputs>(numel, f, data, input_calc, output_calc);
+  } else {
+    auto input_calc = make_input_offset_calculator<num_inputs>(iter);
+    auto output_calc = make_output_offset_calculator<num_outputs>(iter);
+    launch_unrolled_kernel_for_multi_outputs<num_outputs>(numel, f, data, input_calc, output_calc);
+  }
+}
+
+template <typename func_t>
+void gpu_kernel_multiple_outputs(at::TensorIteratorBase& iter, const func_t& f) {
+  for (int arg = 0; arg < iter.ntensors(); arg++) {
+    TORCH_INTERNAL_ASSERT(iter.device(arg).is_xpu());
+  }
+
+  if (iter.numel() == 0) {
+    return;
+  }
+
+  if (!iter.can_use_32bit_indexing()) {
+    for (auto& sub_iter : iter.with_32bit_indexing()) {
+      gpu_kernel_multiple_outputs(sub_iter, f);
+    }
+    return;
+  }
+
+  gpu_kernel_multiple_outputs_impl(iter, f);
+}
\ No newline at end of file
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/tests/__pycache__/__init__.cpython-311.pyc b/tests/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2f7aa6a894229d1bc01abcb0d6773d461aa62f2d
Binary files /dev/null and b/tests/__pycache__/__init__.cpython-311.pyc differ
diff --git a/tests/__pycache__/test_perf.cpython-311-pytest-8.4.1.pyc b/tests/__pycache__/test_perf.cpython-311-pytest-8.4.1.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..44b1fa322bc9fd838343052b7412d5bf17e6e55a
Binary files /dev/null and b/tests/__pycache__/test_perf.cpython-311-pytest-8.4.1.pyc differ
diff --git a/tests/__pycache__/test_rotary.cpython-311-pytest-8.4.1.pyc b/tests/__pycache__/test_rotary.cpython-311-pytest-8.4.1.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fb5321204b00d0bc152284390878b54cdcaf94c4
Binary files /dev/null and b/tests/__pycache__/test_rotary.cpython-311-pytest-8.4.1.pyc differ
diff --git a/tests/__pycache__/test_rotary_sycl.cpython-311-pytest-8.4.1.pyc b/tests/__pycache__/test_rotary_sycl.cpython-311-pytest-8.4.1.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ac513c7e9342a826b4ad18b1b45787acdc5b86b2
Binary files /dev/null and b/tests/__pycache__/test_rotary_sycl.cpython-311-pytest-8.4.1.pyc differ
diff --git a/tests/__pycache__/test_rotary_triton.cpython-311-pytest-8.4.1.pyc b/tests/__pycache__/test_rotary_triton.cpython-311-pytest-8.4.1.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1e9694e6c0885936066b467175559f94de0fe94e
Binary files /dev/null and b/tests/__pycache__/test_rotary_triton.cpython-311-pytest-8.4.1.pyc differ
diff --git a/tests/__pycache__/utils.cpython-311.pyc b/tests/__pycache__/utils.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..42eb5b94702098e8be1e2c4745c740f80602f358
Binary files /dev/null and b/tests/__pycache__/utils.cpython-311.pyc differ
diff --git a/tests/test_rotary.py b/tests/test_rotary.py
new file mode 100644
index 0000000000000000000000000000000000000000..5bcd66105e63a490d17c811ef72531d20e80888f
--- /dev/null
+++ b/tests/test_rotary.py
@@ -0,0 +1,127 @@
+import pytest
+import torch
+
+from tests.utils import infer_device, supports_bfloat16
+from kernels import get_local_kernel
+from pathlib import Path
+
+# import rotary
+# from transformers.trainer_utils import set_seed
+# set_seed(42)
+
+# Set the local repo path, relative path
+repo_path = Path(__file__).parent.parent
+rotary = get_local_kernel(repo_path=repo_path, package_name="rotary")
+
+def apply_rotary_torch(x1: torch.Tensor, x2: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, conj: bool = False):
+    assert x1.shape == x2.shape, "x1 and x2 must have the same shape"
+
+    if not conj:
+        out1 = x1 * cos - x2 * sin
+        out2 = x1 * sin + x2 * cos
+    else:
+        out1 = x1 * cos + x2 * sin
+        out2 = -x1 * sin + x2 * cos
+    return out1, out2
+
+
+def apply_rotary_torch_wrapper(q, k, cos, sin, conj: bool = False):
+    """the wrapper for apply_rotary_torch"""
+    rotary_dim = cos.shape[-1]
+    
+    # apply rotation encoding to Q
+    q1 = q[..., :rotary_dim]
+    q2 = q[..., rotary_dim : 2 * rotary_dim]
+    q_out_1, q_out_2 = apply_rotary_torch(q1, q2, cos, sin, conj)
+    q_out = torch.cat([q_out_1, q_out_2, q[..., 2 * rotary_dim:]], dim=-1)
+
+    # apply rotation encoding to K
+    k1 = k[..., :rotary_dim]
+    k2 = k[..., rotary_dim : 2 * rotary_dim]
+    k_out_1, k_out_2 = apply_rotary_torch(k1, k2, cos, sin, conj)
+    k_out = torch.cat([k_out_1, k_out_2, k[..., 2 * rotary_dim:]], dim=-1)
+
+    return q_out, k_out
+
+
+def apply_rotary_kernel_wrapper(q, k, cos, sin, conj: bool = False):
+    """the wrapper for apply_rotary_kernel"""
+    rotary_dim = cos.shape[-1]
+    
+    # apply rotation encoding to Q
+    q1 = q[..., :rotary_dim]
+    q2 = q[..., rotary_dim : 2 * rotary_dim]
+    rotary.apply_rotary(q1, q2, cos, sin, q1, q2, conj)
+
+    # apply rotation encoding to K
+    k1 = k[..., :rotary_dim]
+    k2 = k[..., rotary_dim : 2 * rotary_dim]
+    rotary.apply_rotary(k1, k2, cos, sin, k1, k2, conj)
+
+
+@pytest.mark.parametrize("batch_size", [1, 2])
+@pytest.mark.parametrize("nheads", [8, 16])
+@pytest.mark.parametrize("seqlen", [128, 256])
+@pytest.mark.parametrize("headdim, rotary_dim", [(64, 32), (128, 64), (64, 30)])
+@pytest.mark.parametrize("qk_dim", [3, 4])
+@pytest.mark.parametrize(
+    "dtype, atol, rtol",
+    [
+        (torch.float32, 1e-5, 1e-5),
+        pytest.param(
+            torch.bfloat16,
+            1e-1,
+            1e-5,
+            marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
+        ),
+    ],
+)
+@pytest.mark.parametrize("conj", [False, True])
+@pytest.mark.flaky(max_runs=2, min_passes=1)
+def test_rotary_equivalence(batch_size, nheads, seqlen, headdim, rotary_dim, qk_dim, dtype, atol, rtol, conj):
+    device = infer_device()
+    if device is None:
+        pytest.skip("No suitable device found for testing")
+
+    if qk_dim == 4:
+        q_shape = (batch_size, seqlen, nheads, headdim)
+        cos_sin_shape = (seqlen, 1, rotary_dim)
+    elif qk_dim == 3:
+        q_shape = (batch_size * seqlen, nheads, headdim)
+        cos_sin_shape = (batch_size * seqlen, 1, rotary_dim)
+
+    q_orig = torch.randn(q_shape, device=device, dtype=dtype)
+    k_orig = torch.randn(q_shape, device=device, dtype=dtype)
+    cos = torch.randn(cos_sin_shape, device=device, dtype=dtype)
+    sin = torch.randn(cos_sin_shape, device=device, dtype=dtype)
+
+    q_kernel, k_kernel = q_orig.clone(), k_orig.clone()
+    q_torch, k_torch = q_orig.clone(), k_orig.clone()
+
+    q_torch_out, k_torch_out = apply_rotary_torch_wrapper(q_torch, k_torch, cos, sin, conj)
+    apply_rotary_kernel_wrapper(q_kernel, k_kernel, cos, sin, conj)
+
+    # verify the rotation results of Q and K are consistent
+    try:
+        assert torch.allclose(q_torch_out, q_kernel, atol=atol, rtol=rtol), "Rotary transformation results for Q do not match"
+    except AssertionError:
+        diff_q = torch.abs(q_torch_out - q_kernel)
+        max_diff_q = torch.max(diff_q)
+        print(f"Max difference for Q: {max_diff_q}")
+        raise
+    try:
+        assert torch.allclose(k_torch_out, k_kernel, atol=atol, rtol=rtol), "Rotary transformation results for K do not match"
+    except AssertionError:
+        diff_k = torch.abs(k_torch_out - k_kernel)
+        max_diff_k = torch.max(diff_k)
+        print(f"Max difference for K: {max_diff_k}")
+        raise
+
+    # verify the non-rotated part of Q and K remains unchanged
+    if (2 * rotary_dim) < headdim:
+        assert torch.equal(
+            q_kernel[..., 2 * rotary_dim:], q_orig[..., 2 * rotary_dim:]
+        ), "Non-rotated part of Q should be unchanged"
+        assert torch.equal(
+            k_kernel[..., 2 * rotary_dim:], k_orig[..., 2 * rotary_dim:]
+        ), "Non-rotated part of K should be unchanged"
diff --git a/tests/utils.py b/tests/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..049b7ce0868d5751d812df47c412ca5a5c11aab9
--- /dev/null
+++ b/tests/utils.py
@@ -0,0 +1,23 @@
+import torch
+
+
+def infer_device():
+    """
+    Get current device name based on available devices
+    """
+    if torch.cuda.is_available():  # Works for both Nvidia and AMD
+        return "cuda"
+    elif torch.xpu.is_available():
+        return "xpu"
+    else:
+        return None
+
+
+def supports_bfloat16():
+    device = infer_device()
+    if device == "cuda":
+        return torch.cuda.get_device_capability() >= (8, 0)  # Ampere and newer
+    elif device == "xpu":
+        return True
+    else:
+        return False
\ No newline at end of file
diff --git a/torch-ext/torch_binding.cpp b/torch-ext/torch_binding.cpp
index 22086369b4ba06ca1896424ab51d38b6ec1a3d44..51f71d2062c03fcb29885767a621b44020d924c7 100644
--- a/torch-ext/torch_binding.cpp
+++ b/torch-ext/torch_binding.cpp
@@ -1,12 +1,17 @@
 #include <torch/all.h>
+
+#if defined(CUDA_KERNEL)
 #include <c10/cuda/CUDAGuard.h>
+#elif defined(XPU_KERNEL)
+#include <c10/core/DeviceGuard.h>
+#endif
 
 #include "registration.h"
 
-#define CHECK_DEVICE(x) TORCH_CHECK(x.device().type() == torch::kCUDA, #x " must be on CUDA")
+#define CHECK_DEVICE(x) TORCH_CHECK(x.device().type() == torch::kCUDA || x.device().type() == torch::kXPU, #x " must be on CUDA or XPU")
 #define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
 
-void apply_rotary_cuda(torch::Tensor const &x1, torch::Tensor const &x2,
+void _apply_rotary(torch::Tensor const &x1, torch::Tensor const &x2,
                        torch::Tensor const &cos, torch::Tensor const &sin,
                        torch::Tensor &out1, torch::Tensor &out2,
                        bool const conj);
@@ -27,16 +32,23 @@ void apply_rotary(torch::Tensor const &x1, torch::Tensor const &x2,
     TORCH_CHECK(cos.sizes() == sin.sizes());
     TORCH_CHECK(out1.sizes() == out2.sizes());
 
+#if defined(CUDA_KERNEL)
     // Otherwise the kernel will be launched from cuda:0 device
     at::cuda::CUDAGuard device_guard{x1.device()};
-
-    apply_rotary_cuda(x1, x2, cos, sin, out1, out2, conj);
+#elif defined(XPU_KERNEL)
+    c10::DeviceGuard device_guard{x1.device()};
+#endif
+    _apply_rotary(x1, x2, cos, sin, out1, out2, conj);
 }
 
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("apply_rotary(Tensor x1, Tensor x2, Tensor cos, Tensor sin,"
           "Tensor! out1, Tensor! out2, bool conj) -> ()");
-  ops.impl("apply_rotary", torch::kCUDA, &apply_rotary);
+#if defined(CUDA_KERNEL)
+    ops.impl("apply_rotary", torch::kCUDA, &apply_rotary);
+#elif defined(XPU_KERNEL)
+    ops.impl("apply_rotary", torch::kXPU, &apply_rotary);
+#endif
 }
 
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)