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

fix: handle zero-strided tensors in fast_rope_embedding (#3781) [skip pre-commit.ci]

Browse source
This commit is contained in:
ゆり 2026-03-12 08:42:47 +00:00 • committed by Daniel Han
commit 3d64e11cc7
1 changed files with 6 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

12
unsloth/kernels/rope_embedding.py
View file

@ -311,9 +311,9 @@ class Fast_RoPE_Embedding_QK(torch.autograd.Function):
batch, n_heads_Q, seq_len, head_dim = Q.shape
_, 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
# Clone if not contiguous or has zero strides, such as expanded tensors.
Q_out = Q.clone() if not Q.is_contiguous() or 0 in Q.stride() else Q
K_out = K.clone() if not K.is_contiguous() or 0 in K.stride() else K
if has_indices:
# TRL's rotary indices are always in int32, so casting is just for safety
@ -383,9 +383,9 @@ 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
# Clone if not contiguous or has zero strides, such as expanded tensors.
dQ_out = dQ.clone() if not dQ.is_contiguous() or 0 in dQ.stride() else dQ
dK_out = dK.clone() if not dK.is_contiguous() or 0 in dK.stride() else dK
Q_batch_stride, Q_head_stride, Q_seq_stride = (
dQ_out.stride(0),