don't save inputs/outputs buffer of FlashAttenFunc to reduce memory usage for inference mode

[XiaobingSuper]

In inference mode, we don't need to save the inputs/outputs used by the training mode. This PR can reduce memory usage during LLM serving(such as vllm, it uses those flash attn APIs for a better performance).

[rocking5566]

@XiaobingSuper
I found this PR make pytest (https://github.com/Dao-AILab/flash-attention/tree/main/tests) fail.

flash-attention/flash_attn/flash_attn_interface.py

Line 455 in 0dfb281

is_grad = torch.is_grad_enabled() and qkv.requires_grad

torch.is_grad_enabled() always be False in the test script, which make backward cannot access ctx.saved_tensors

[XiaobingSuper]

@XiaobingSuper I found this PR make pytest (https://github.com/Dao-AILab/flash-attention/tree/main/tests) fail.

flash-attention/flash_attn/flash_attn_interface.py

Line 455 in 0dfb281

is_grad = torch.is_grad_enabled() and qkv.requires_grad

torch.is_grad_enabled() always be False in the test script, which make backward cannot access ctx.saved_tensors

Sorry, it is my mistake, gradient computation is already disabled in custom autograd.Functions by default(https://discuss.pytorch.org/t/is-torch-no-grad-making-a-difference-in-custom-autograd-functions/186627), it need to check it before calling the customer op.

[XiaobingSuper]

@rocking5566 I created a PR #1397 to fix this issue. Thanks.