From 78a5e58932b47a34748aa66f46609f2d30b818b1 Mon Sep 17 00:00:00 2001 From: Lei Zhenyuan Date: Thu, 10 Jul 2025 05:10:25 +0800 Subject: [PATCH] [bugs] fix for casual mask (#2868) * fix for casual mask * use un_casual in sdpa * add missing mask * fix for type --- unsloth/models/llama.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index f08b4762eb..ca33509f64 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -324,6 +324,13 @@ def LlamaAttention_fast_forward_inference( # Knn, Vnn = Knn, Vnn # 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 # Attention if bsz == 1: Qn *= self.scalar # See https://github.com/ggerganov/llama.cpp/issues/7805#issuecomment-2153349963 @@ -524,11 +531,18 @@ def LlamaAttention_fast_forward( V = V.transpose(1, 2) A = flash_attn_func(Q, K, V, causal = True) else: + # 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 # Grouped query attention if SDPA_HAS_GQA: # 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, enable_gqa = n_groups != 1) + A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = is_causal, enable_gqa = n_groups != 1) # Go back to (batch_size, seq_len, n_heads, head_dim) A = A.transpose(1, 2)#.contiguous() else: @@ -543,7 +557,7 @@ def LlamaAttention_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) + 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