[Paddle] Some refactor and optimization on fused attention

[Wong4j]

1. Get rng state in c++ custom op and extract seed/offset using a CUDA kernel. (similar to PyTorch)
2. Add two fused CUDA kernel to convert mask -> actual seqlens -> cu seqlens.
3. Add more fused_attention_backend check in python code.

[Wong4j]

@jeng1220 Please take a look.

[Wong4j]

This PR will have some conflicts with PR#410.
Please merge PR#410 first. Then I will fix the conflicts.

[Wong4j]

@jeng1220
Changes in the latest 2 commits:

• Use CUB to do prefix sum
• Use paddle.incubate.nn.layer.fused_dropout_add

[mingxu1067]

/te-ci

[zlsh80826]

/te-ci

[jeng1220]

LGTM.

[jeng1220]

The PR #412 needs to be merged first.

[mingxu1067]

/te-ci

[jeng1220]

@timmoon10 ,

There are some UT failures but they are not related to this PR:

• Build / PyTorch (pull_request)
• Lint / PyTorch Python (pull_request)

Please review and merge the code if everything looks good.
Thank you.

[timmoon10]

LGTM, sorry for taking such a long time.

Be advised that #403 is adding a more flexible C API for fused attention, although the existing API is still supported.

[jeng1220]

LGTM, sorry for taking such a long time.

Be advised that #403 is adding a more flexible C API for fused attention, although the existing API is still supported.

It seems that the new #403 changes the order of output dimensions of fused attention.

- # output: [b, s, h]+ # output: [s, b, h]

Here: https://github.com/NVIDIA/TransformerEngine/pull/403/files#diff-8646f274b30b8b4b8c13804b7df0d5032eb53cd033a4fff11aa91053a14f7206R642

It will take some time to study how to integrate it into TE/Paddle.
@timmoon10 , I hope this PR can be merged first.

[timmoon10]

That comment was incorrect, attention has already been using [s, b, h, 3, d] tensors. I see that Paddle expects [b, s, 3, h, d] tensors, so we've found a bug. The Paddle backend transposes the tensors:

TransformerEngine/transformer_engine/paddle/layer/attention.py

Lines 341 to 343 in 3f64de4

	q=paddle.transpose(x=q, perm=[0, 2, 1, 3])
	k=paddle.transpose(x=k, perm=[0, 2, 1, 3])
	v=paddle.transpose(x=v, perm=[0, 2, 1, 3])

Could we do something similar for the TE backend?

[Wong4j]

Thank you @timmoon10 . I also noticed TE-PyT use sbh3d layout by default.
TE-Paddle only supports bs3hd layout for backend=='transformer_engine' currently. As for backend=='paddle', I simply split the fused qkv with shape bs3hd into three bshd tensor, and transpose them to bhsd, and do the dot product attention calculation with paddle native op. We have verified in UT that the calculation results of the two backends are consistent. It is no bug.
I will continue to improve this feature to support more qkv_layout in next PR, which will involve some big changes to the code. We can merge this PR first.

[timmoon10]

/te-ci

[timmoon10]

/te-ci

Edit: Having trouble launching the job. See pipeline 10073671.