AttnFuncWithCP can use less memory

In AttnFuncWithCP.forward, up to 3 buffer (two calculating and one recving) is used simultaneously but p2p_comm_buffers will grow to cp_size which I believe is a waste. Just add buffer[i] = None and kv_input[i%2] = None after _flash_attn_forward may lower memory usage as the reference count of calculated kv is zero.

Here's some mock code:

• No buffer release:

import torch

torch.cuda.memory._record_memory_history(max_entries=100000)

device = 'cuda:7'
cp_size = 16
streams = [torch.cuda.current_stream(), torch.cuda.Stream()]
buffers = [None for _ in range(cp_size)]
buffers[0] = torch.empty(100, 100, dtype=torch.float32, device=device)

torch.cuda.memory._dump_snapshot('dev_cuda_mem_1.pkl')
for i in range(cp_size):
    with torch.cuda.stream(streams[i % 2]):
        if i < cp_size - 1:
            buffers[i + 1] = torch.empty(100, 100, dtype=torch.float32, device=device)
            torch.cuda.memory._dump_snapshot('cuda_mem_snapshot.pkl')

        kv = buffers[i]
        # ....
        for _ in range(10000):
            buffers[i] *= buffers[i]

        # buffers[i] = None

• With buffer release:

import torch

torch.cuda.memory._record_memory_history(max_entries=100000)

device = 'cuda:7'
cp_size = 16
streams = [torch.cuda.current_stream(), torch.cuda.Stream()]
buffers = [None for _ in range(cp_size)]
buffers[0] = torch.empty(100, 100, dtype=torch.float32, device=device)

torch.cuda.memory._dump_snapshot('dev_cuda_mem_1.pkl')
for i in range(cp_size):
    with torch.cuda.stream(streams[i % 2]):
        if i < cp_size - 1:
            buffers[i + 1] = torch.empty(100, 100, dtype=torch.float32, device=device)
            torch.cuda.memory._dump_snapshot('cuda_mem_snapshot.pkl')

        kv = buffers[i]
        # ....
        for _ in range(10000):
            buffers[i] *= buffers[i]

        buffers[i] = None

Link to PR #951

Start at AttnFuncWithCP.forward and compare the linked PR #951 with the issue's mock buffer-release examples. Check memory behavior around _flash_attn_forward, p2p_comm_buffers, and cp_size; done means the buffers can be released without changing attention behavior or causing communication errors.