From 1279a6168171c859d70d46727975eb541de4fa55 Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Tue, 13 Jan 2026 17:46:54 +0800 Subject: [PATCH] [update] prompt prefill --- model/model_minimind.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/model/model_minimind.py b/model/model_minimind.py index 830c78d..d826f82 100755 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -179,7 +179,7 @@ class Attention(nn.Module): xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) cos, sin = position_embeddings - xq, xk = apply_rotary_pos_emb(xq, xk, cos[:seq_len], sin[:seq_len]) + xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin) # kv_cache实现 if past_key_value is not None: @@ -193,14 +193,11 @@ class Attention(nn.Module): repeat_kv(xv, self.n_rep).transpose(1, 2) ) - if self.flash and seq_len > 1 and (attention_mask is None or torch.all(attention_mask == 1)): + if self.flash and (seq_len > 1) and (past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)): output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=True) else: scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim) - scores = scores + torch.triu( - torch.full((seq_len, seq_len), float("-inf"), device=scores.device), - diagonal=1 - ).unsqueeze(0).unsqueeze(0) # scores+mask + scores[:, :, :, -seq_len:] += torch.triu(torch.full((seq_len, seq_len), float("-inf"), device=scores.device), diagonal=1) if attention_mask is not None: extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)