From 3d64e11cc7b4bfa70a5fd28711f7a3924900bf41 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E3=82=86=E3=82=8A?=
 <249254018+yurekami@users.noreply.github.com>
Date: Thu, 12 Mar 2026 08:42:47 +0000
Subject: [PATCH] fix: handle zero-strided tensors in fast_rope_embedding
 (#3781) [skip pre-commit.ci]

---
 unsloth/kernels/rope_embedding.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/unsloth/kernels/rope_embedding.py b/unsloth/kernels/rope_embedding.py
index fcc9cb923b..18dfb00373 100644
--- a/unsloth/kernels/rope_embedding.py
+++ b/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),