Add unit tests for shared prefix masked attention with torch.FlexAttention

[austin362667]

Summary

TLDR of #476: The shared prefix attention mask is an optimization for paired-preference alignment training.

To pave the way for #476, this PR aims to set up basic unit tests of flex attn with casual and shared prefix mask.

Testing Done

Benchmarks

1. Casual Attention Mask (Flash Attention 2 vs. Torch Scaled Dot Product Attention vs. FlexAttention)

3. Shared Prefix Attention Mask (Flash Attention 2 vs. Torch Scaled Dot Product Attention vs. FlexAttention)

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[austin362667]

The FlexAttention correctness check failed only on AMD GPU might be related to this ISSUE and PR.