diff --git a/src/h3_blackwell_runtime/attention.py b/src/h3_blackwell_runtime/attention.py
index f8362ee..0878ef5 100644
--- a/src/h3_blackwell_runtime/attention.py
+++ b/src/h3_blackwell_runtime/attention.py
@@ -86,15 +86,12 @@ class H3SageAttention(nn.Module):
         sequence = x.shape[0]
         inner = self.heads * self.head_dim
         q, k, v = self.qkv_proj(x).split(inner, dim=-1)
-        q = q.view(1, sequence, self.heads, self.head_dim)
-        k = k.view(1, sequence, self.heads, self.head_dim)
+        q = rms_norm(q.view(1, sequence, self.heads, self.head_dim), self.q_norm_weight, self.eps)
+        k = rms_norm(k.view(1, sequence, self.heads, self.head_dim), self.k_norm_weight, self.eps)
         v = v.view(1, sequence, self.heads, self.head_dim)
 
-        from comfy_kitchen import rms_rope_split_half_
-
-        rms_rope_split_half_(q, k, rope_rotation, self.q_norm_weight, self.k_norm_weight, epsilon=self.eps, rot_dim=rope_rotation.shape[-3] * 2)
-        q = q.transpose(1, 2).contiguous()
-        k = k.transpose(1, 2).contiguous()
+        q = apply_split_half_rope(q, rope_rotation).transpose(1, 2).contiguous()
+        k = apply_split_half_rope(k, rope_rotation).transpose(1, 2).contiguous()
         v = v.transpose(1, 2).contiguous()
 
         if self.backend == "sage2":