Skip to content

[dev] Reuse parent communicators for dynamic context parallelism - #6801

Draft
xiaoyao0115 wants to merge 1 commit into
NVIDIA:devfrom
xiaoyao0115:dcp-cg-parent-communicators
Draft

[dev] Reuse parent communicators for dynamic context parallelism#6801
xiaoyao0115 wants to merge 1 commit into
NVIDIA:devfrom
xiaoyao0115:dcp-cg-parent-communicators

Conversation

@xiaoyao0115

@xiaoyao0115 xiaoyao0115 commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

Summary

  • keep logical dynamic-CP groups for rank/topology, but reuse the existing parent DP×CP communicator for Transformer Engine P2P attention
  • route dynamic-CP router subgroup reductions through fixed, disjoint lanes on the same parent communicator
  • avoid eagerly initializing a separate NCCL communicator for every logical CP-size group when the parent transport is available

The logical groups still determine CP size, rank, and peers. Only the communication transport is reused, so non-P2P and local implementations keep their existing behavior.

Implementation

MCore registers each logical CP group with its parent DP×CP group before TE attention. TE selects peers from the logical group but issues P2P on the parent ProcessGroup. Router reductions use a parent-sized fixed layout so different dynamic subgroups do not overlap.

This works independently of CUDA Graph. CUDA Graph slot/memory sharing is the follow-up in #6479.

Dependencies and merge order

  1. [PyTorch] Allow CP P2P transport group overrides TransformerEngine#3420
  2. this PR
  3. [dev] Share CUDA graph slots across dynamic CP variants #6479

This PR targets dev; #6479 is stacked on pull-request/6801.

Validation

  • Black and isort checks
  • 11 targeted communicator/router unit tests
  • 4-GPU dynamic-subgroup reduction tests
  • combined 4-GPU CUDA Graph probe: 50 replays with mixed CP2/CP1/CP1

@copy-pr-bot

copy-pr-bot Bot commented Aug 24, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@xiaoyao0115

Copy link
Copy Markdown
Contributor Author

/ok to test d03753d

@xiaoyao0115

Copy link
Copy Markdown
Contributor Author

/ok to test cd8e9b2

@xiaoyao0115
xiaoyao0115 force-pushed the dcp-cg-parent-communicators branch from cd8e9b2 to d7cb2de Compare August 25, 2026 09:49
@xiaoyao0115 xiaoyao0115 changed the title [dev] Reuse parent communicators for dynamic CP CUDA graphs [dev] Reuse parent communicators for dynamic context parallelism Aug 25, 2026
@xiaoyao0115
xiaoyao0115 changed the base branch from pull-request/6479 to dev August 25, 2026 09:49

Copy link
Copy Markdown
Contributor Author

/ok to test d7cb2de

Signed-off-by: Tailai Ma <tailaim@ferries-work.nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant