From ac7831126170e538a8dcbfde5f84f95f8591e6eb Mon Sep 17 00:00:00 2001 From: Lei Zhenyuan Date: Tue, 26 Aug 2025 19:45:20 +0800 Subject: [PATCH] fix is casual for qwen3 (#3213) --- unsloth/models/qwen3.py | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py index b20f22dab1..95d7232002 100644 --- a/unsloth/models/qwen3.py +++ b/unsloth/models/qwen3.py @@ -185,7 +185,15 @@ def Qwen3Attention_fast_forward( Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous() # Needs (batch_size, n_heads, seq_len, head_dim) # is_casual and attention_mask must not be both set! - A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False) + # when qlen==vlen and attn_mask is None, we should use causal attention + Q_len = Q.shape[-2] + K_len = K.shape[-2] + if attention_mask is None and Q_len == K_len: + is_causal = True + else: + is_causal = False + + A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = is_causal) # Go back to (batch_size, seq_len, n_heads, head_dim) A = A.transpose(1, 2).contiguous() pass @@ -336,6 +344,14 @@ def Qwen3Attention_fast_forward_inference( Knn, Vnn = Kn, Vn pass + # when qlen==vlen and attn_mask is None, we should use causal attention + Q_len = Qn.shape[-2] + K_len = Knn.shape[-2] + if attention_mask is None and Q_len == K_len: + is_causal = True + else: + is_causal = False + # Grouped query attention _, _, cached_len, _ = Knn.shape if bsz == 1 or not SDPA_HAS_GQA and n_groups != 1: @@ -358,9 +374,9 @@ def Qwen3Attention_fast_forward_inference( A = torch_matmul(A, Vnn, out = Qn) else: if SDPA_HAS_GQA: - A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False, enable_gqa = True) + A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = is_causal, enable_gqa = True) else: - A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False) + A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = is_causal) pass A = A.transpose(1, 2) A = A.reshape(bsz, 1, attention_size)