Skip to content

GitLab

Commit c55de02f authored by Azure's avatar Azure
Browse files

fix qlen > 1000 mask is none error

parent 8747c099
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
ktransformers/operators/attention.py
View file @ c55de02f
...@@ -195,11 +195,11 @@ class KDeepseekV2Attention(BaseInjectedModule, DeepseekV2Attention): ...@@ -195,11 +195,11 @@ class KDeepseekV2Attention(BaseInjectedModule, DeepseekV2Attention):
[:,:min(self.chunck_size, min(past_key_value.max_cache_len-cur_idx, self.chunck_size))] [:,:min(self.chunck_size, min(past_key_value.max_cache_len-cur_idx, self.chunck_size))]
self.attn_mask[:, :, :, cur_idx+self.chunck_size:] = -1e+38 self.attn_mask[:, :, :, cur_idx+self.chunck_size:] = -1e+38
self.attn_mask[:, :, :, :cur_idx] = 0 self.attn_mask[:, :, :, :cur_idx] = 0
chunck_mask = torch.narrow(self.attn_mask, 2, 0, min(self.chunck_size, q_len-cur_idx)) chunk_mask = torch.narrow(self.attn_mask, 2, 0, min(self.chunck_size, q_len-cur_idx))
cur_output, _, _ = self.forward_chunck( cur_output, _, _ = self.forward_chunck(
hidden_states[:, cur_idx:min(cur_idx + self.chunck_size, q_len), ...], hidden_states[:, cur_idx:min(cur_idx + self.chunck_size, q_len), ...],
chunck_mask, chunk_mask,
position_ids[:, cur_idx:min(cur_idx + self.chunck_size, q_len)], position_ids[:, cur_idx:min(cur_idx + self.chunck_size, q_len)],
past_key_value, past_key_value,
output_attentions, output_attentions,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment