Implement improved attention masking for bidirectional_masked

llama_diffusion_model.py

@@ -47,10 +47,16 @@ class BidirectionalLlamaAttention(LlamaAttention):
         key_states = self.repeat_kv(key, module.num_key_value_groups)
         value_states = self.repeat_kv(value, module.num_key_value_groups)
 
-        attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+        # attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+        # if attention_mask is not None:
+        #     causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        #     attn_weights = attn_weights + causal_mask
+
         if attention_mask is not None:
-            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
-            attn_weights = attn_weights + causal_mask
+            # Convert bool -> float with -inf where masked
+            attn_mask = attention_mask.masked_fill(~attention_mask, float('-inf')).to(query.dtype)
+            attn_weights = attn_weights + attn_mask
+
 
         attn_weights = nn.functional.softmax(attn_weights, dim=-1).to(query.dtype)
         attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)