ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[Intel] make intel device support ROPE (#3164)

Browse source 
* make intel device pass

* abstract torch device stream
This commit is contained in:
Lei Zhenyuan 2025-09-03 19:39:57 +08:00 • committed by GitHub
commit 5f14e6fcd4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 5 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

7
unsloth/kernels/rope_embedding.py
View file

@ -15,7 +15,7 @@
import triton
import triton.language as tl
import torch
from .utils import calculate_settings, torch_gpu_device
from .utils import calculate_settings, torch_gpu_device, torch_device_stream
ROPE_GROUP_SIZE : int = 4
def _rope_embedding(
@ -156,8 +156,7 @@ 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()
torch_device_stream(Q.device).synchronize()
return Q, K
pass
@ -201,6 +200,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()
torch_device_stream(Q.device).synchronize()
return Q, K
pass

2
unsloth/kernels/utils.py
View file

@ -174,6 +174,8 @@ else:
cgemm_4bit_inference_naive_bf16 = bnb.functional.lib.cgemm_4bit_inference_naive_bf16
pass
torch_device_stream = torch.xpu.current_stream if DEVICE_TYPE == "xpu" else torch.cuda.current_stream
torch_mm = torch.mm
torch_mv = torch.mv
torch_matmul = torch.matmul