From b52bd92e9a87a913352617fd5864f0a027a66c3d Mon Sep 17 00:00:00 2001 From: Francesco Bertolotti Date: Fri, 21 Aug 2026 12:08:01 +0200 Subject: [PATCH] free unused kvs Signed-off-by: Francesco Bertolotti --- .../attention/dot_product_attention/context_parallel.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index ea89ca97eb..0ca6513574 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -1715,6 +1715,12 @@ def forward( ) kv_inputs[i % 2] = p2p_comm_buffers[i] + + # Release the KV chunk from two steps back. + if i >= 2 and not is_graph_capturing(): + p2p_comm_buffers[i - 2].record_stream(flash_attn_streams[i % 2]) + p2p_comm_buffers[i - 2] = None + k_part = kv_inputs[i % 2][:k_numel].view(*k_shape) v_part = kv_inputs[i % 2][k_numel:].view(*v_shape) q_part = q