[Jetteyc]

Description

• 补充变量解决NameError: name 'nvshmem_kv' is not defined和NameError: name 'cp_global_ranks' is not defined报错。
• 当cp和其他并行同时使用时会报错，但是单独使用不会报错。调用顺序如下：context_parallel_nvshmem.py里transformer_engine_torch编译自transformer_engine/pytorch/csrc，tex.nvshmem_get_on_current_stream调用的是transformer_engine/pytorch/csrc/extensions/nvshmem_comm.cpp，最终用的是transformer_engine::pytorch的nvshmemx_getmem_on_stream。在context_parallel_nvshmem.py的tex.init_nvshmem_backend(cp_group)传的是cp的编号，所以tex.nvshmem_get_on_current_stream中最后一个参数peer应该传的是cp rank而不是全局rank。context_parallel_nvshmem.py的三处tex.nvshmem_get_on_current_stream中最后一个参数把int(owner_global)改成int(owner_idx)

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

• transformer_engine/pytorch/attention/dot_product_attention/context_parallel_nvshmem.py: 补充nvshmem_kv和cp_global_ranks定义
• transformer_engine/pytorch/attention/dot_product_attention/context_parallel_nvshmem.py: 三处tex.nvshmem_get_on_current_stream最后一个参数把int(owner_global)改成int(owner_idx)

Checklist:

• I have read and followed the contributing guidelines
• 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

[Copilot]

Pull request overview

This PR addresses runtime errors and incorrect NVSHMEM peer selection in the PyTorch context-parallel attention implementation that uses NVSHMEM for KV exchange.

Changes:

• Persist nvshmem_kv on the autograd context so it is available in backward().
• Switch NVSHMEM peer argument from a global rank mapping to the CP-group rank (owner_idx) for correct addressing when CP is combined with other parallelism modes.
• Remove the now-unused global-rank mapping usage around NVSHMEM gets.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The NVSHMEM fallback path is effectively unreachable here: init_nvshmem_backend() / create_nvshmem_tensor() will raise if NVSHMEM isn't available, but the try/except is commented out and later code still checks if nvshmem_kv is not None to decide between NVSHMEM vs P2P. Either restore the guarded fallback (so nvshmem_kv can legitimately be None) or remove the if nvshmem_kv is not None branches and related commented code to keep behavior and comments consistent.

[Copilot]

cp_global_ranks = ctx.cp_global_ranks is now unused in backward() after switching NVSHMEM peer addressing to owner_idx. Please remove this assignment (or use it) to avoid dead code and confusion about whether peer should be global rank or process-group rank.

Suggested change

cp_global_ranks = ctx.cp_global_ranks