[JAX] Consolidate the distributed fused attention test code (#1405)

[JAX] Consolidate the distributed fused attention test code (#1405) #66