Add more unit tests to FA fwd kernels.

[xinyazhang]

Note it is not testing the backward kernel but use the kernel in ref_bwd*.py as reference.

Run the UT with pytest test_backward.py

To run a know set of parameters, change the main2 function and run python test_backward.py (pytest -k works but it's much slower and requires -s to enable standard output)

The whole test suite takes around 12 hours to complete (after disabling auto-tuning). The main problem is the Triton kernel compiling. With more tl.constexpr it make take even longer.

[xinyazhang]

Known problem: for Triton commit 00e09cf3008b86978f25f838659698e4a0bf6f45. Running pytest test_backward.py -v -x shows the following runtime error.

self = <.HIPLauncher object at 0x79492bd4c220>, args = (1, 1, 1, 180074496, 185729744, (8, 1, 32768, 1, 1, 1), ...), kwargs = {}

    def __call__(self, *args, **kwargs):
        print(f'{args=}')
        print(f'{kwargs=}')
>       self.launch(*args, **kwargs)
E       RuntimeError: Triton Error [HIP]:  Code: 1, Messsage: invalid argument

../../../aotriton/third_party/triton/python/triton/backends/amd/driver.py:420: RuntimeError

Removing all autotune configs (except for 'BLOCK_M': 16, 'BLOCK_N': 16,) can mitigate this problem, but this is probably not what we want.

[xinyazhang]

Known problem: for Triton commit 00e09cf3008b86978f25f838659698e4a0bf6f45. Running pytest test_backward.py -v -x shows the following runtime error.

Confirmed this is caused by double loading of libamdhip64.so and can be fixed by a189c11

As a temporary solution, setting TRITON_LIBHIP_PATH to PyTorch's .so file can fix this as well