From ccefc2e337306d1bc45ab38165530c1adefa56ad Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Sat, 2 Aug 2025 16:00:34 +0530 Subject: [PATCH] fixup rope sync for everything (#3061) --- unsloth/kernels/rope_embedding.py | 6 +++++- unsloth/models/llama.py | 2 -- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/unsloth/kernels/rope_embedding.py b/unsloth/kernels/rope_embedding.py index a06d81a512..46b6766edb 100644 --- a/unsloth/kernels/rope_embedding.py +++ b/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 diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index e7d9084ad8..3c0d5012ae 100644 --- a/unsloth/models/llama.py +++ b/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)