add zero_tensors option with default False to flash_attn_varlen_func api

[wangye805]

Motivation

TransformerEngine (TE) integrate flash-attn as one of the attn backends. For the varlen mode in TE, users can create q, k, v with a fixed total_seqlen larger than the cu_seqlen[-1]. Therefore the padding area for O, dQ, dK, dV, from [cu_seqlen[-1]: end] is undefined since we allocate those tensors with torch.empty_like(). Undefined tensors in the padding area can affect the subsequent bwd operations in transformer model. In addition, on NV h100, torch.empty_like already returns zeroed-out tensors.

Therefore we propose an extra zero_tensors option with default value False into api flash_attn_varlen_func. For customers who creates q/k/v with fixed length, we pass this option into ck kernels so that they can zero out padding areas in the output. This won't affect the original default behavior

Technical Details

Add the zero_tensors option into flash_attn_varlen_func api

Test Plan

Will test flash-attn CI and TE CI

Test Result

TE CI passed
Flash-attn CI pending

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[rocking5566]

Could you send PR to upstream directly? tridao branch (here) only sync fork with upstream

[wangye805]

Will send PR to upstream directly