can mask be supported?

[unwritten]

in
flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False):

can mask be supported before going into softmax like:

        att = att.masked_fill(mask[:, None, :, :], min_value)
        att = F.softmax(att, dim=-1, dtype=torch.float32).to(rdtype).masked_fill(mask[:, None, :, :], 0.0)

[adityas-amd]

Hi @unwritten,
I think it is supported on main branch,

flash-attention/flash_attn/flash_attn_interface.py

Lines 77 to 106 in 99730ad

	def _flash_attn_forward(
	q: torch.Tensor,
	k: torch.Tensor,
	v: torch.Tensor,
	dropout_p: float,
	softmax_scale: float,
	causal: bool,
	window_size_left: int,
	window_size_right: int,
	softcap: float,
	alibi_slopes: Optional[torch.Tensor],
	return_softmax: bool
	) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
	q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
	out, softmax_lse, S_dmask, rng_state = flash_attn_cuda.fwd(
	q,
	k,
	v,
	None,
	alibi_slopes,
	dropout_p,
	softmax_scale,
	causal,
	window_size_left,
	window_size_right,
	softcap,
	return_softmax,
	None,
	)
	return out, softmax_lse, S_dmask, rng_state