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

fix is casual for qwen3 (#3213)

Browse source
This commit is contained in:
Lei Zhenyuan 2025-08-26 19:45:20 +08:00 • committed by GitHub
commit ac78311261
1 changed files with 19 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

22
unsloth/models/qwen3.py
View file

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