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

fastrope fix for zero strided tensors (#3782)

Browse source 
Co-authored-by: Francesco Bertolotti <francesco.bertolotti@igenius.ai>
This commit is contained in:
Francesco Bertolotti 2025-12-29 06:21:48 +01:00 • committed by GitHub
commit dabf2a901b
1 changed files with 12 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

24
unsloth/kernels/rope_embedding.py
View file

@ -312,8 +312,8 @@ class Fast_RoPE_Embedding_QK(torch.autograd.Function):
_, n_heads_K, _, _ = K.shape
# Inplace rotary embedding is generally fine
Q_out = Q.clone() if not Q.is_contiguous else Q
K_out = K.clone() if not K.is_contiguous else K
Q_out = Q.clone() if not Q.is_contiguous() else Q
K_out = K.clone() if not K.is_contiguous() else K
if has_indices:
# TRL's rotary indices are always in int32, so casting is just for safety
@ -383,21 +383,21 @@ class Fast_RoPE_Embedding_QK(torch.autograd.Function):
else ctx.cos.new_empty(1, dtype = torch.int32)
)
# Inplace rotary embedding is generally fine
dQ_out = dQ.clone() if not dQ.is_contiguous() else dQ
dK_out = dK.clone() if not dK.is_contiguous() else dK
Q_batch_stride, Q_head_stride, Q_seq_stride = (
dQ.stride(0),
dQ.stride(1),
dQ.stride(2),
dQ_out.stride(0),
dQ_out.stride(1),
dQ_out.stride(2),
)
K_batch_stride, K_head_stride, K_seq_stride = (
dK.stride(0),
dK.stride(1),
dK.stride(2),
dK_out.stride(0),
dK_out.stride(1),
dK_out.stride(2),
)
# Inplace rotary embedding is generally fine
dQ_out = dQ.clone() if not dQ.is_contiguous else dQ
dK_out = dK.clone() if not dK.is_contiguous else dK
with torch_gpu_device(dQ.device):
_rope_embedding_QK[(batch * ctx.seq_len, ctx.n_heads_Q)](
dQ_out,