diff --git a/unsloth/kernels/rope_embedding.py b/unsloth/kernels/rope_embedding.py
index e93cbd1544..2adc9ecc5a 100644
--- a/unsloth/kernels/rope_embedding.py
+++ b/unsloth/kernels/rope_embedding.py
@@ -20,13 +20,6 @@ from ..device_type import DEVICE_COUNT
 from .utils import calculate_settings, torch_gpu_device, torch_device_stream
 
 
-@triton.heuristics(
-    {
-        "BACKWARD_PASS": lambda args: bool(args["BACKWARD_PASS"]),
-        "HAS_ROPE_INDICES": lambda args: bool(args["HAS_ROPE_INDICES"]),
-    }
-)
-@triton.jit
 def _rope_embedding_QK(
     Q,
     Q_batch_stride,
@@ -104,9 +97,17 @@ def _rope_embedding_QK(
         tl.store(k_ptr + half_head_dim + col_offsets, k1 * cos1 + k0 * sin1, mask = mask)
 
 
-ROPE_GROUP_SIZE: int = 4
+_rope_embedding_QK = triton.jit(_rope_embedding_QK)
+_rope_embedding_QK = triton.heuristics(
+    {
+        "BACKWARD_PASS": lambda args: bool(args["BACKWARD_PASS"]),
+        "HAS_ROPE_INDICES": lambda args: bool(args["HAS_ROPE_INDICES"]),
+    }
+)(_rope_embedding_QK)
 
 
+ROPE_GROUP_SIZE: int = 4
+
 def _rope_embedding(
     Q,
     Q_row_stride: tl.constexpr,