Summary of Changes

Hello @wangfakang, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces symmetric memory allocation for buffers involved in cp-attention-allgather operations. By wrapping tensor creation calls with a use_symmetric_memory context manager, the change aims to optimize memory usage and potentially improve performance in distributed attention mechanisms, ensuring memory is allocated efficiently across devices when symmetric allocation is enabled.

Highlights

• Symmetric Memory Integration: The pull request integrates symmetric memory allocation for cp-attention-allgather buffers, specifically for input_tensor_all and output_tensor.
• Context Manager Application: Memory allocations for key tensors within cp_attn_tp_all_gather_reorganazied_into_tensor and cp_all_gather_rerange_output functions are now managed using the use_symmetric_memory context manager.
• Conditional Allocation: The symmetric memory allocation is applied conditionally, based on the result of the is_allocation_symmetric() check, allowing for flexible memory management.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Friendly ping @merrymercy @Fridge003 Could you help me review this? Thanks!

/tag-and-rerun-ci

After reviewing the failed CI logs, most errors appear to be related to missing python3 command, which doesn't seem to be caused by this PR. How should we proceed with this? cc @nvcastet @ishandhanani @merrymercy @Fridge003

+ echo CUDA_VISIBLE_DEVICES=
CUDA_VISIBLE_DEVICES=
+ python3 -c 'import os, shutil, tempfile, getpass; cache_dir = os.environ.get("TORCHINDUCTOR_CACHE_DIR") or os.path.join(tempfile.gettempdir(), "torchinductor_" + getpass.getuser()); shutil.rmtree(cache_dir, ignore_errors=True)'
scripts/ci/cuda/ci_install_dependency.sh: line 27: python3: command not found
Error: Process completed with exit code 127.

@ishandhanani @Fridge003 I agree with @wangfakang, the CI failures are unrelated to this PR.

@wangfakang In fact, can you add back disabled=not is_allocation_symmetric() to the context manager as a safety net in your PR?
We have seen other places that did not have it fail with misalignement.

/tag-and-rerun-ci

@nvcastet @merrymercy @Fridge003 @ShangmingCai
After reviewing some of the failed CI runs, the failures appear to be caused by flaky test cases. Moreover, the same test cases are also failing in PR#18408 that have already merged into the main branch, such as stage-b-test-small-1-gpu-amd (linux-mi325-gpu-1, 10).

CC: @xu-yfei

LGTM

/rerun-stage stage-c-test-8-gpu-h200

✅ Triggered stage-c-test-8-gpu-h200 to run independently (skipping dependencies).

🔗 View workflow run

/rerun-stage stage-c-test-8-gpu-h20

✅ Triggered stage-c-test-8-gpu-h20 to run independently (skipping dependencies).

🔗 View workflow run

/rerun-stage unit-test-backend-4-gpu-b200

✅ Triggered unit-test-backend-4-gpu-b200 to run independently (skipping dependencies).

🔗 View workflow run

/rerun-stage stage-c-test-4-gpu-b200

✅ Triggered stage-c-test-4-gpu-b200 to run independently (skipping dependencies).

🔗 View workflow run

/rerun-stage stage-c-test-4-gpu-gb200

✅ Triggered stage-c-test-4-gpu-gb200 to run independently (skipping dependencies).

🔗 View workflow run

Since our CI is flaky currently, I think we can bypass CI and merge this PR. If any bug has been reported, please ping me to fix or revert.