Performant backward Triton implementation with separated dkdv and dq kernels

[jtang10]

This PR introduces a performant version of the backward kernel implementation in Triton. It follows the same implementation strategy as the example in upstream Triton but with added functionality.

As expected, it surpasses the existing implementation but falls short to only half of the performance from the example mentioned above. Here's a quick summary of the performance at the time of submitting this PR.

Note that the tutorials/06-fused-attention.py assumes Q and KV have the same heads and sequence length, so it only is tested on those examples.

The performance investigation will follow through from another PR later

[jtang10]

The switch between this new kernel bwd_prefill_split.py and bwd_prefill.py is controlled by an environment variable here

flash-attention/flash_attn/flash_attn_triton_amd/utils.py

Line 14 in ea973eb

USE_SINGLE_BWD_KERNEL = os.environ.get('USE_SINGLE_BWD_KERNEL', '0').lower() in ('1', 'true', 'yes')

.

Though we discussed that an argument in a function to toggle it is preferred, passing the triton kernel wrapper into interfaces makes it hard to be controlled by argument imo, so I leave it as-is, but turn on split kernel by default.

[jtang10]

most of the changes related to removing/changing the existing debug messages is within a single commit and therefore dropped.

[jtang10]

The problem introduced from 58941ed has been manually reverted from 8436dc7 and fixed in 814471e.

The solution is simple, just do not assume any tensor sharing the same stride and therefore use the strides from the tensor itself. For example, q and dq does not necessarily share the same stride because q is individually potentially padded (padding seems to pack the tensor and make it contiguous first) and loaded from forward() path and dq is indexed from dqkv in the backward() path and discontiguous on the dim of 3 at (b, s, 3, h, d). As a result, assuming q and dq would lead to wrong results.

Another change is to NOT .contiguous() q, k, v at the beginning of triton wrapper. After the change above, this is unnecessary and create additional data movement.

TODO:
I keep the hack from your version removed but didn't patch this update into it. I'm gonna focus on the perf part so I'll leave it to you @micmelesse to have it updated.

[micmelesse]

This is excellent