Flash-Attention-Softmax-N
Flash-Attention-Softmax-N copied to clipboard
added unit tests for `attn_mask`
I added unit tests for the case when the attn_mask arugment of flash_attention_n or slow_attention_n is not None in an attempt to reproduce #39. My unit tests pass on my personal machine (cpu).