remove contiguous copy for flash-attn opbuilder

[YizhouZ]

following megatron-ds original attention shape, now flash-attn op builder supports only [B, S, H, D] shape. Two contiguous copy time could be improved by this PR.

[YizhouZ]

@conglongli Hi, could you help to review this PR? Thanks

[YizhouZ]

Hi @tjruwase, it seems that conglong is not available, could you help to review this? Thanks.

[tjruwase]

@YizhouZ, will this change work with older flash-attn op builder?

[YizhouZ]

@YizhouZ, will this change work with older flash-attn op builder?

I think not. These are two different shape required for input tensor.

I could make it specific for xpu device if you want less impact.

[tjruwase]

I think not. These are two different shape required for input tensor.

In that case, can we condition based on flash attn version? It is important to maintain backwards-compatibility. Thanks!

[YizhouZ]

I think not. These are two different shape required for input tensor.

In that case, can we condition based on flash attn version? It is important to maintain backwards-compatibility. Thanks!

I made some changes:

• isolate flash_attn_builder from other flash attn usages
• compatibility is ensured by checking function name, new version would be flash_attn_builder.flash_attn_func_v2

could you help to review again? Thanks!

[tjruwase]

@YizhouZ, thanks for the PR.