Add varlen support to AOTriton's Flash Attention

[xinyazhang]

Varlen Flash Attention is implemented by two new APIs: attn_fwd_compact_varlen and attn_bwd_compact_varlen, with the same set of kernels. Check include/aotriton/flash.h for their details.

Note:

1. the bias tensor input b is reserved. Uses should pass TensorView<4>::get_null_tensor() for this argument for now. Any other inputs are not supported nor tested.
2. The varlen API still expects Rank-4 tensors, which uniforms the code b/w varlen and non-varlen. The API also expects the size of the first dimension (batch) is exact 1. Extra batches will not be processed. This interface is slightly different from Tri Dao's implementation (https://github.com/Dao-AILab/flash-attention/blob/320fb59487658f033f56711efd3d61b7c7a6f8f3/csrc/flash_attn/flash_api.cpp#L500-L502), and users are supposed to use torch.transpose and torch.unsqueeze to match the API.
3. This PR also refactored the Triton kernel, and unified the compiled data type of sequence length related arguments to int32_t

[wangye805]

I took a quick look and realized your implementation should work for the qkv_layout of THD in NVTE.

How about for the layout of BSHD/BHSD with padding mask and cu_seqlen, as we discussed before?

[xinyazhang]

How about for the layout of BSHD/BHSD with padding mask and cu_seqlen, as we discussed before?

It's in the triton kernel but not exposed as C++ API. They will be added in a separate PR as attn_?wd_padded_varlen

[wangye805]

How about for the layout of BSHD/BHSD with padding mask and cu_seqlen, as we discussed before?

It's in the triton kernel but not exposed as C++ API. They will be added in a separate PR as attn_?wd_padded_varlen

Do you mind pointing me to the triton source codes for padded_varlen?

[xinyazhang]

Do you mind pointing me to the triton source codes for padded_varlen?

aotriton/tritonsrc/fwd_kernel.py

Lines 73 to 85 in c8551b1

	else: # < 0 for padded seqlen
	cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z)
	cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1)
	seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start
	if start_m * BLOCK_M >= seqlen_q:
	return
	cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z)
	cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1)
	seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start
	# Varlen, but padded to Rank 4 tensor
	cu_seqlens_q_start = 0
	cu_seqlens_k_start = 0
	batch_index = off_z

[wangye805]

LGTM

[groenenboomj]

Looking to add more support for more layouts later but that will be a different piece of work. This looks to meet needs for varlen.

[groenenboomj]

Fine for this PR but will we be adding some layout flag or defaulting to this for other layouts?

[xinyazhang]

What do you mean by layout flag? The current plan is to let all inputs be BHSD layout.
If new layouts are needed we are going to add new APIs instead of changing existing ones.