[attention] Make the FlexAttention mask type configurable - #5

[attention] Make the FlexAttention mask type configurable#5

aws-arashj wants to merge 2 commits into

mainaws-neuron/torchtitan-neuron:mainfrom

attention/flex-mask-typeaws-neuron/torchtitan-neuron:attention/flex-mask-typeCopy head branch name to clipboard