test: refactor flash_attn tests to use parameterized (#20913)
Use `parameterized` to decompose the huge test case. This will make
adding ROCm support be possible.
---------
Co-authored-by: Guangyun Han <guangyunhan@microsoft.com@h100vm-ort.kxelwkzfzxguje5bxvwxxs135a.gvxx.internal.cloudapp.net>