vrr/kvcache-ai-ktransformers
2
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2025-09-07 21:19:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix qlen > 1000 mask is none error

Browse source
This commit is contained in:
Azure 2024-09-02 02:58:10 +00:00
parent 8747c099f2
commit c55de02f7b
1 changed files with 2 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
ktransformers/operators/attention.py
View file

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