[PyTorch] Add no-load-balance THD all-gather CP by sudhakarsingh27 · Pull Request #3221 · NVIDIA/TransformerEngine

sudhakarsingh27 · 2026-07-18T00:35:32Z

How to use

Select the token partition strategy explicitly and use the same value for both
context-parallel attention configuration and input partitioning:

from transformer_engine.pytorch import CPAttentionLoadBalancingStrategy

strategy = CPAttentionLoadBalancingStrategy.NO_LOAD_BALANCE

layer.set_context_parallel_group(
    cp_group,
    cp_global_ranks,
    cp_stream,
    cp_comm_type="all_gather",
    load_balancing_strategy=strategy,
)

input_ids, labels, position_ids = get_batch_on_this_cp_rank(
    cu_seqlens_padded,
    input_ids,
    labels,
    position_ids,
    cp_group=cp_group,
    load_balancing_strategy=strategy,
)

CPAttentionLoadBalancingStrategy.DUAL_CHUNK_SWAP remains the default.

What changed

Add the public CPAttentionLoadBalancingStrategy enum with
DUAL_CHUNK_SWAP and experimental NO_LOAD_BALANCE strategies.
Propagate the selected strategy through TransformerLayer,
MultiheadAttention, DotProductAttention, and the context-parallel attention
backends.
Assign one contiguous physical-token chunk to each CP rank for
NO_LOAD_BALANCE, while preserving logical document boundaries through THD
metadata.
Capture the selected strategy for backward so forward and backward use the
same token layout.
Reuse the existing native THD partition-index implementation for CUDA batch
slicing while retaining the CPU dataloader fallback.
Add focused utility, propagation, FusedAttention, and gated unpadded
FlashAttention 3 coverage.

Why

The default per-document DualChunkSwap partition divides every sequence into
2 * cp_size chunks and performs two attention steps per rank. Partitioning
the complete physical buffer into cp_size contiguous chunks allows one
attention step per rank and supports packed documents whose individual lengths
are not divisible by 2 * cp_size.

This strategy intentionally trades causal load balance for fewer attention
calls.

Scope and constraints

The experimental strategy currently requires THD self-attention,
cp_comm_type="all_gather", causal attention without a sliding window
(window_size=(-1, 0)), equal local Q/K/V physical lengths, and either
FusedAttention or FlashAttention 3. FlashAttention 3 requires
pad_between_seqs=False. FP8 and CUDA graph capture are not supported.

Input partitioning and attention configuration must use the same strategy.

Validation

21 focused context-parallel utility tests passed.
CP2 BF16 THD all-gather FusedAttention forward/backward passed for both
DUAL_CHUNK_SWAP and NO_LOAD_BALANCE.
TransformerLayer strategy propagation and the default
DUAL_CHUNK_SWAP behavior passed focused smoke checks.
Autograd forward/backward arity validation and git diff --check passed.
Pylint passed with a 10.00/10 score on the changed production files.

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>

greptile-apps · 2026-08-18T23:08:40Z

Greptile Summary

The PR adds an explicit context-parallel load-balancing strategy and implements contiguous physical-token partitioning for no-load-balance THD all-gather attention.

Propagates the strategy through TransformerLayer, MultiheadAttention, DotProductAttention, and attention backends.
Adds strategy-aware THD partitioning, causal metadata generation, K/V ordering, and backward-state capture.
Adds focused utility and distributed FusedAttention/FlashAttention 3 coverage.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the previously reported backward-layout instability is addressed by consistently using the strategy captured during forward.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Implements contiguous THD partitioning and consistently captures the selected strategy for matching forward and backward layouts.
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Stores and propagates the context-parallel load-balancing strategy into supported attention backends.
transformer_engine/pytorch/attention/multi_head_attention.py	Propagates context-parallel strategy configuration to core attention.
transformer_engine/pytorch/transformer.py	Exposes strategy configuration through the TransformerLayer composition.
tests/pytorch/attention/test_cp_utils.py	Covers strategy-aware partitioning, metadata generation, and stable K/V restoration behavior.
tests/pytorch/attention/test_attention_with_cp.py	Adds distributed no-load-balance coverage for FusedAttention and supported FlashAttention 3 configurations.

Sequence Diagram

sequenceDiagram
    participant Input as Input partitioner
    participant Layer as TransformerLayer / MHA
    participant DPA as DotProductAttention
    participant CP as THD all-gather CP
    participant Bwd as Autograd backward
    Input->>Input: Partition tokens using selected strategy
    Layer->>DPA: Propagate strategy
    DPA->>CP: Forward Q/K/V and strategy
    CP->>CP: Gather K/V and build strategy-specific metadata
    CP->>CP: Save strategy on autograd context
    CP-->>Layer: Attention output
    Bwd->>CP: Restore saved strategy
    CP->>CP: Select matching step layout and K/V ordering
    CP-->>Bwd: dQ/dK/dV

_{Reviews (7): Last reviewed commit: "Skip known sm90 deterministic THD OOM" | Re-trigger Greptile}

Name the single contiguous-chunk policy after its deliberate lack of causal load balancing so the performance tradeoff is explicit. Remove the CPU reference partitioner and require the existing CUDA path for per-document metadata. Capture the selected layout through backward and add focused helper plus CP2 forward/backward coverage so mutable environment state cannot make the two passes use different token orders. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

Validate the experimental policy where context-parallel communication is selected so unsupported combinations fail before entering custom autograd. Passing the captured mode into the internal all-gather call also prevents a second environment lookup from selecting a different layout. Restore the default CPU dataloader slicing behavior, keep one stream dependency per format path, and exercise padded feature execution so the lean experimental path does not regress existing callers. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

cyanguwa · 2026-08-19T19:52:40Z


 # Float8CurrentScaling: fused_attn_bwd takes O in FP8 by default, this flag allows it in F16
 _dpa_fp8_cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1"
+_NO_LOAD_BALANCE_ENV = "NVTE_EXPERIMENTAL_CP_AG_THD_NO_LOAD_BALANCE"


Jax CP defines "DualChunkSwap", "Stripped", and "NoLoadBalance" as some Strategy (an enum). Maybe we can do the same on the PyTorch side? We can still label it experimental that way (in the docs), but without the environment variable.

Done. The environment flag is replaced by the public CPLoadBalancingStrategy enum. NO_LOAD_BALANCE is documented as experimental, and DUAL_CHUNK_SWAP remains the default.

cyanguwa · 2026-08-20T19:18:01Z

        fp8_meta,
        quantizers,
        fp8_output,
+        no_load_balance,


These "no_load_balance" flags in code and tests can all become an Enum, or a flag for load_balancing_strategy I think.

Done. The booleans are replaced with CPLoadBalancingStrategy, propagated through the attention stack, and captured for backward.

cyanguwa · 2026-08-20T19:19:26Z

-        cp_stream.wait_stream(torch.cuda.current_stream())
-
+            # Preserve overlap by letting cp_stream proceed before output initialization.
+            cp_stream.wait_stream(torch.cuda.current_stream())


Why shifting wait_stream to only the else branch? Any particular reason for this?

The THD wait follows K/V restoration and metadata preparation; the non-THD wait remains after K/V reordering. Comments now clarify both dependencies.

cyanguwa · 2026-08-20T19:21:30Z

+    contiguous physical-buffer chunk to each rank and uses one attention step per rank.
+    Logical sequences remain isolated by ``cu_seqlens``. This mode requires THD,
+    all-gather, full causal self-attention, and FusedAttention, or FlashAttention 3
+    without padding. Input producers must use the same flag when partitioning inputs


I guess "without padding" here means pad_between_seqs=False?

Yes—this specifically means pad_between_seqs=False. The documentation and validation now state that explicitly.

cyanguwa · 2026-08-20T19:24:00Z

+        assert "causal" in attn_mask_type and window_size == (
+            -1,
+            0,
+        ), "No-load-balance THD partitioning currently supports full causal attention only."


I think you mentioned somewhere else about "full causal attention". I understand you mean "no sliding window", but I feel generally when people say causal mask or causal attention, they do mean "full causal"? (nit)

Updated to “causal attention without a sliding window,” specifically window_size=(-1, 0).

cyanguwa · 2026-08-20T19:26:33Z

+                def build_rank_indices(device):
+                    # Preserve the CPU-capable per-document dataloader path.
+                    rank_slices = []
+                    for slice_size, seq_start in zip(slice_sizes, cu_seqlens_padded[:-1]):


Is there any clever way to do this without this loop? Can you ask your agent? :)

Done. CUDA partitioning now uses tex.thd_get_partitioned_indices; the Python loop remains only as the CPU fallback.

Refresh the PR on current upstream before addressing review feedback. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

Make token partitioning explicit so input slicing and attention cannot diverge through mutable process state. Reuse native THD indices for CUDA while preserving CPU dataloader behavior, and cover the supported FusedAttention and unpadded FlashAttention 3 paths. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

Use the shorter public name because the strategy governs both input partitioning and attention execution. Consolidate redundant low-level partition tests into the existing end-to-end slicing coverage. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

Document the experimental contract at the public selection points. Preserve the legacy child-setter invocation for the default strategy so the additive API does not disrupt existing extension modules. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

sudhakarsingh27 · 2026-08-21T08:49:18Z

/te-ci pytorch L1

The focused no-load-balance test uses cp_2_0, which reaches the same known cuDNN deterministic THD backward workspace limit as the generic CP matrix. Apply the equivalent sm90 skip so the standalone coverage does not bypass that guard. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

sudhakarsingh27 · 2026-08-21T20:41:16Z

/te-ci pytorch

sudhakarsingh27 force-pushed the sudhakars/relaxed-load-balancing-all-gather branch from 50b86ff to 9645f28 Compare July 18, 2026 00:39

sudhakarsingh27 requested review from KshitijLakhani and cyanguwa July 18, 2026 00:43

sudhakarsingh27 added 6 commits August 18, 2026 14:40

sudhakarsingh27 force-pushed the sudhakars/relaxed-load-balancing-all-gather branch from 9645f28 to 57f3cf8 Compare August 18, 2026 21:44

sudhakarsingh27 marked this pull request as ready for review August 18, 2026 23:02

sudhakarsingh27 self-assigned this Aug 18, 2026

sudhakarsingh27 added the 2.19 label Aug 18, 2026

greptile-apps Bot reviewed Aug 18, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py Outdated

sudhakarsingh27 changed the title ~~[PyTorch] Add experimental packed-contiguous THD all-gather CP~~ [PyTorch] Add no-load-balance THD all-gather CP Aug 19, 2026

sudhakarsingh27 commented Aug 20, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py Outdated

cyanguwa reviewed Aug 20, 2026

View reviewed changes

sudhakarsingh27 added 2 commits August 20, 2026 23:52

Merge origin/main into no-load-balance THD branch

eb543ee

Refresh the PR on current upstream before addressing review feedback. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

sudhakarsingh27 requested a review from ksivaman as a code owner August 21, 2026 08:09

sudhakarsingh27 added 2 commits August 21, 2026 01:26

cyanguwa previously approved these changes Aug 21, 2026

View reviewed changes

sudhakarsingh27 dismissed cyanguwa’s stale review via ac516e8 August 21, 2026 18:59

Uh oh!

Conversation

sudhakarsingh27 commented Jul 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

How to use

What changed

Why

Scope and constraints

Validation

Uh oh!

greptile-apps Bot commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Sequence Diagram

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sudhakarsingh27 commented Aug 21, 2026

Uh oh!

sudhakarsingh27 commented Aug 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

sudhakarsingh27 commented Jul 18, 2026 •

edited

Loading

greptile-apps Bot commented Aug 18, 2026 •

edited

Loading