refactor flash attn test (#21028)

Commit

1 year ago

refactor flash attn test (#21028) Code improvement and allow change to use other ep