fixed masked flash attention

OpenNMT · May 17, 2024 · 2c4eded · 2c4eded
1 parent 0920080
commit 2c4eded
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/onmt/modules/multi_headed_attn.py b/onmt/modules/multi_headed_attn.py
@@ -483,6 +483,12 @@ def forward(
                         if sliding_window > 0 and key.size(2) > sliding_window:
                             key = key[:, :, 1:, :]
                             value = value[:, :, 1:, :]
+                    if key_pad_mask is not None and step == 0:
+                        x = key_pad_mask
+                        x = x.expand(-1, self.head_count // self.parallel_gpu, -1)
+                        x = x.unsqueeze(3)
+                        x = x.expand(-1, -1, -1, 128)
+                        value = value.masked_fill(x, 0)
 
                     self.layer_cache[1]["keys"] = key
                     self.layer_cache[1]["values"] = value