[PyTorch] Add no-load-balance THD all-gather CP

[sudhakarsingh27]

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.

[greptile-apps]

Greptile Summary

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

• Propagates the strategy from TransformerLayer through the attention modules and backend adapters.
• Captures the forward strategy for consistent backward token restoration and step selection.
• Adds THD metadata, partitioning helpers, validation, and distributed backend coverage.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Adds contiguous THD partitioning and consistently uses the forward-captured strategy for backward restoration, step selection, and gradient reduce-scatter layout.
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Stores and forwards the selected context-parallel strategy through backend dispatch, including checkpointed attention calls.
transformer_engine/pytorch/attention/multi_head_attention.py	Propagates context-parallel strategy configuration to core dot-product attention.
transformer_engine/pytorch/transformer.py	Exposes strategy configuration at the TransformerLayer boundary and forwards it to attention children.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Passes the strategy into FlashAttention and FusedAttention context-parallel execution.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Updates attention utility behavior to account for the selected context-parallel partition strategy.
transformer_engine/pytorch/constants.py	Defines the public context-parallel load-balancing strategy enum.
transformer_engine/pytorch/init.py	Re-exports the load-balancing strategy through the public PyTorch package.
tests/pytorch/attention/test_cp_utils.py	Adds focused coverage for contiguous partition metadata, captured restoration mode, and CPU/CUDA rank slicing.
tests/pytorch/attention/test_attention_with_cp.py	Adds distributed FusedAttention and FlashAttention 3 coverage for no-load-balance THD attention.
tests/pytorch/attention/run_attention_with_cp.py	Extends the distributed attention harness to generate and partition inputs according to the selected strategy.

Sequence Diagram

sequenceDiagram
    participant U as Caller
    participant TL as TransformerLayer
    participant MHA as MultiheadAttention
    participant DPA as DotProductAttention
    participant CP as CP All-Gather Attention
    U->>TL: set_context_parallel_group(strategy)
    TL->>MHA: propagate CP configuration
    MHA->>DPA: propagate strategy
    DPA->>CP: forward(Q, K, V, strategy)
    CP->>CP: partition and build THD metadata
    CP->>CP: save strategy on autograd context
    CP-->>DPA: attention output
    Note over CP: Backward uses captured strategy
    CP->>CP: restore K/V and select step layout
    CP->>CP: unrestore dK/dV for reduce-scatter

Loading

_{Reviews (10): Last reviewed commit: "Limit deterministic THD guard to F16 bac..." | Re-trigger Greptile}

[sudhakarsingh27]

/te-ci pytorch L1

[sudhakarsingh27]

/te-ci pytorch

[cyanguwa]

Two minor comments - please run the local CI at least and I'm happy to approve again. Thanks for the PR!