Attention fwd/bwd issues with all zero masks

Hi,

We have been observing some issues with masking, namely that if all tokens are masked out. In this specific case, the output of attention will be twice of what native pytorch's implementation is, and the gradient of v value vector is simply garbage value.

To demonstrate the need for masking all tokens, consider the case where we land on no templates in template search but the computation graph needs to stay the same during training to avoid a bunch of issues including memory fragmentation.

Is there a way to fix this?

Best,
Rui

Start by reproducing the attention forward and backward pass with an all-zero mask, then compare its output and the gradient of the value vector with native PyTorch. The issue does not name a source file or test, so locate the attention masking implementation and its existing tests first. Done means the all-masked case has matching forward results and a valid value gradient while preserving the computation graph.