Update transformer.py

transformer.py

@@ -189,10 +189,8 @@ class GroupedQueryAttention(nn.Module):
                 if attention_mask.dim() == 2:
                     attention_mask = attention_mask[:, None, None, :]
                 if attention_mask.dtype != torch.float:
-                    # 假设传入的是 1(Keep)/0(Mask)
                     extended_mask = (1.0 - attention_mask) * torch.finfo(attn_scores.dtype).min
                 else:
-                    # 假设传入的已经是加性 mask (0/-inf)
                     extended_mask = attention_mask
             
                 attn_scores = attn_scores + extended_mask
@@ -203,7 +201,7 @@ class GroupedQueryAttention(nn.Module):
                 torch.ones(seq_len_k, seq_len_k, device=x.device, dtype=torch.bool),
                 diagonal=1
                 )
-                causal_mask = causal_mask[-q.shape[2]:, :]#还没懂
+                causal_mask = causal_mask[-q.shape[2]:, :]
                 attn_scores = attn_scores.masked_fill(causal_mask, float('-inf'))
             
             attention_weights = F.softmax(attn_scores, dim=-1, dtype=torch.float32).to(q.dtype)