Ring Attention: free unused kv comm buffers

[francesco-bertolotti]

Description

In Ring Attention, p2p_comm_buffers[i] holds KV chunks received from other ranks. As the ring rotates, each rank receives KV chunks from its peers, and p2p_comm_buffers grows to eventually hold the full KV cache. Depending on the context length, head size, and data type, this can consume several GB of GPU memory.

To put this into perspective, consider a sequence of 1M tokens with a head size of 256, 8 KV heads, and bfloat16. The KV cache requires approximately:

1e6 × 256 × 8 × 2 × 2B = 8 GB

This can significantly reduce the available GPU memory by the end of Ring Attention.

However, the buffer appears to be larger than necessary. At time step i, a rank only needs:

• p2p_comm_buffers[i], which is used to run FlashAttention and send the current KV chunk to the next rank.
• p2p_comm_buffers[i+1], which is used to receive the next KV chunk.

There are no references to p2p_comm_buffers[j] for j < i. Therefore, once a KV chunk has been processed and forwarded, it can be deallocated. At any point, only the current and next KV chunks need to remain allocated.

I ran the relevant tests on A100 GPUs, which are the hardware currently available to me, and all tests passed.

At extreme context lengths, this three-line change saves several GB of GPU memory and enables Ring Attention to scale to 1M-token contexts. Without this change, we hit an OOM.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Deallocate unused KV chunks from p2p_comm_buffers during Ring Attention.

Checklist

• I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst)
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR reduces Ring Attention GPU-memory growth by releasing processed KV communication chunks after their final use.

• Records the relevant FlashAttention stream before dropping each old buffer reference.
• Preserves the current and next KV chunks needed by the rotating P2P ring.
• Leaves CUDA graph capture behavior unchanged by disabling eviction while graph capture is active.

Confidence Score: 4/5

The PR appears safe to merge, with a non-blocking request to add coverage for the newly introduced KV-buffer eviction path at context-parallel sizes of at least three.

The released buffer has no subsequent forward or backward use, its attention consumer runs on the recorded stream, and outstanding P2P work is waited before release; only targeted regression coverage is missing.

Files Needing Attention: transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Adds stream-aware release of processed Ring Attention KV buffers; the lifetime logic appears sound, but the newly active multi-rank eviction branch lacks direct regression coverage.

_{Reviews (1): Last reviewed commit: "free unused kvs" | Re-trigger Greptile}

[greptile-apps]

Cover the active eviction branch

The new buffer release executes only for pure P2P Ring Attention with a context-parallel size of at least three, but existing coverage does not exercise that configuration. Add a regression test for this branch so stream-lifetime and ring-index regressions do not leave larger runs vulnerable to corrupted results or GPU memory errors.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!