[CP] Register KV cache allgather buffer with symmetric memory (#24040)

Signed-off-by: wangfakang <fakangwang@gmail.com>
2026-05-06 23:24:36 +08:00
parent c4c5541618
commit bfc1aeae13
1 changed files with 9 additions and 6 deletions
@@ -179,6 +179,9 @@ def cp_all_gather_reorganized_into_tensor_kv_cache(
        input_tensor = F.pad(input_tensor, padding, mode="constant", value=0)

    # Create output tensor with proper shape for all dimensions
+    with use_symmetric_memory(
+        get_attention_cp_group(), disabled=not is_allocation_symmetric()
+    ):
        input_tensor_full = torch.empty(
            max_len * cp_size,
            *input_tensor.shape[1:],