fixup rope sync for everything (#3061)

unsloth/kernels/rope_embedding.py

@@ -93,7 +93,7 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
         # [TODO] Changing blocksize to head_dim//2 seems to have
         # some concurrency / un-deterministic issues.
         BLOCK_SIZE, num_warps = calculate_settings(head_dim//2) # (head_dim//2)
-        
+
         # group_size = 4 # 4 or 8, too large group_size can hurt performance.
         div : int
         mod : int
@@ -155,6 +155,9 @@ pass
 def fast_rope_embedding(Q, K, cos, sin):
     Q = Fast_RoPE_Embedding.apply(Q.transpose(1, 2), cos, sin).transpose(1, 2)
     K = Fast_RoPE_Embedding.apply(K.transpose(1, 2), cos, sin).transpose(1, 2)
+    # synchronize before cat to avoid race condition
+    torch.cuda.current_stream(Q.device).synchronize()
+
     return Q, K
 pass
 
@@ -198,5 +201,6 @@ pass
 def inplace_rope_embedding(Q, K, cos, sin, position_ids):
     Q = Slow_RoPE_Embedding.apply(Q, cos, sin, position_ids)
     K = Slow_RoPE_Embedding.apply(K, cos, sin, position_ids)
+    torch.cuda.current_stream(Q.device).synchronize()
     return Q, K
 pass

unsloth/models/llama.py

@@ -499,8 +499,6 @@ def LlamaAttention_fast_forward(
     #     else inplace_rope_embedding(Q, K, cos, sin, position_ids)
     # )
     Q, K = fast_rope_embedding(Q, K, cos, sin)
-    # synchronize before cat to avoid race condition
-    torch.cuda.current_stream(Q.device).synchronize()
 
     if past_key_value is not None:
         K = torch.cat([past_key_value[0], K], dim = 2)