[PyTorch] Add experimental packed-contiguous THD all-gather CP#3221
Draft
sudhakarsingh27 wants to merge 6 commits into
Draft
[PyTorch] Add experimental packed-contiguous THD all-gather CP#3221sudhakarsingh27 wants to merge 6 commits into
sudhakarsingh27 wants to merge 6 commits into
Conversation
Allow causal THD attention to shard the complete packed token buffer with mirrored context-parallel chunks while retaining document boundaries through sequence metadata. This supports workloads whose individual documents are not divisible by twice the CP size. Keep the existing per-document partition as the default and reject backend or attention combinations that the prototype has not validated, so existing paths remain unchanged. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Reuse the existing THD CUDA partition and reorder kernels by representing the complete physical token buffer as one partitioning sequence. Preserve CPU and mixed-device dataloader behavior with a reference fallback. Rename the opt-in policy to packed_super_sequence to distinguish physical partitioning from THD packing, and add a narrowly gated matched-input benchmark path so the two policies can be compared without workload or timing asymmetry. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Allow relaxed THD all-gather to assign one contiguous chunk per CP rank while preserving the mirrored policy for compatibility and comparison. Rank-major ownership needs no KV reorder and reduces each rank to one attention step; keep it opt-in while the performance tradeoffs are evaluated. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Keep per-document partitioning as the default and packed-contiguous as the only packed opt-in so the experimental API has one global ownership contract. Delete the unused 2*CP global metadata and reorder paths, simplify packed metadata to one chunk and one attention step per rank, and retain a negative test for the retired selector. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Keep the experimental packed-contiguous selection internal to context parallelism so public TransformerLayer and attention APIs retain their existing signatures. Per-document partitioning remains the default unless NVTE_EXPERIMENTAL_CP_AG_THD_PACKED_CONTIGUOUS=1 is set. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Keep the initial upstream review focused on the production context-parallel implementation. The test development remains available in the preceding commits while the branch tip restores the existing test suite unchanged. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
sudhakarsingh27
force-pushed
the
sudhakars/relaxed-load-balancing-all-gather
branch
from
July 18, 2026 00:39
50b86ff to
9645f28
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
How to use
Set
NVTE_EXPERIMENTAL_CP_AG_THD_PACKED_CONTIGUOUS=1to turn the switch internallyWhat changed
NVTE_EXPERIMENTAL_CP_AG_THD_PACKED_CONTIGUOUS=1; the current per-document DualChunkSwap path remains the default.main.Why
The current per-document THD partition divides every sequence into
2 * cp_sizechunks and performs two attention steps per rank. For uneven packed documents, using the complete physical buffer as the sharding unit makes it possible to usecp_sizeequal chunks and one attention step per rank, reducing per-step overhead while retaining document isolation.Scope and constraints
The experimental path currently requires THD self-attention, all-gather CP, a full causal mask, equal local Q/K/V physical lengths, and either FusedAttention or FlashAttention 3 without inter-sequence padding. FP8 and CUDA graph capture are intentionally rejected. Input partitioning and attention must both see the same environment flag.
Validation
py_compile,git diff --check, and autograd forward/backward arity validation passed on the rebased tip.