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

unsloth/kernels/rope_embedding.py

@@ -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),