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

tests/pytorch/attention/run_attention_with_cp.py

@@ -11,16 +11,17 @@
 import torch.distributed as dist
 from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
     get_cu_seqlens_on_cp_rank,
+    get_thd_partitioned_indices,
 )
 from transformer_engine.pytorch.attention.dot_product_attention.utils import combine_and_quantize
-import transformer_engine_torch as tex
 from transformer_engine.pytorch import DType
 from test_attention_with_cp import (
     model_configs_flash_attn,
     model_configs_fused_attn,
 )
 from transformer_engine.pytorch import (
     autocast,
+    CPLoadBalancingStrategy,
     DotProductAttention,
     Float8Quantizer,
     Float8CurrentScalingQuantizer,
@@ -52,6 +53,7 @@ def generate_input_shapes(
     world_size: int,
     kernel_backend: str,
     fa_pad_between_seqs: str = "False",
+    load_balancing_strategy=CPLoadBalancingStrategy.DUAL_CHUNK_SWAP,
 ):
     if qkv_format == "bshd":
         q_input_shape = (
@@ -110,8 +112,33 @@ def generate_input_shapes(
         cu_seqlens_q_padded = None
         cu_seqlens_kv_padded = None
     elif qkv_format == "thd":
-        seqlens_q = torch.randint(0, config.max_seqlen_q + 1, [config.batch_size]).to(torch.int32)
-        seqlens_q_padded = (seqlens_q + 2 * world_size - 1) // (world_size * 2) * (world_size * 2)
+        if load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE:
+            assert config.batch_size == 2
+            if kernel_backend == "FlashAttention" and fa_pad_between_seqs == "False":
+                seqlens_q = torch.tensor(
+                    [config.max_seqlen_q - 2, config.max_seqlen_q], dtype=torch.int32
+                )
+                assert seqlens_q.sum().item() % world_size == 0
+                seqlens_q_padded = seqlens_q
+            else:
+                # Exercise both document padding and a CP chunk boundary inside a document.
+                seqlens_q = torch.tensor(
+                    [config.max_seqlen_q - 2, config.max_seqlen_q - 1],
+                    dtype=torch.int32,
+                )
+                padded_total = 2 * config.max_seqlen_q
+                assert padded_total % world_size == 0
+                seqlens_q_padded = torch.tensor(
+                    [config.max_seqlen_q - 1, config.max_seqlen_q + 1],
+                    dtype=torch.int32,
+                )
+        else:
+            seqlens_q = torch.randint(0, config.max_seqlen_q + 1, [config.batch_size]).to(
+                torch.int32
+            )
+            seqlens_q_padded = (
+                (seqlens_q + 2 * world_size - 1) // (world_size * 2) * (world_size * 2)
+            )
         cu_seqlens_q_padded = torch.cat(
             [
                 torch.zeros([1], dtype=torch.int32),
@@ -206,10 +233,12 @@ def run_dpa_with_cp(
     is_training="True",
     fa_pad_between_seqs="False",
     deterministic="False",
+    load_balancing_strategy="DUAL_CHUNK_SWAP",
     log_level=logging.WARNING,
 ):
     """Test DotProductAttention module with context parallelism"""
     logging.root.setLevel(log_level)
+    load_balancing_strategy = CPLoadBalancingStrategy[load_balancing_strategy]
     # When is_training is False, gradient outputs are None.
     is_training = is_training == "True"
     pad_between_seqs = None
@@ -331,7 +360,14 @@ def run_dpa_with_cp(
         cu_seqlens_kv,
         cu_seqlens_q_padded,
         cu_seqlens_kv_padded,
-    ) = generate_input_shapes(qkv_format, config, world_size, kernel_backend, fa_pad_between_seqs)
+    ) = generate_input_shapes(
+        qkv_format,
+        config,
+        world_size,
+        kernel_backend,
+        fa_pad_between_seqs,
+        load_balancing_strategy,
+    )
     q_orig = torch.clamp(torch.randn(q_input_shape, dtype=dtypes[dtype]), min=-1, max=1).cuda()
     k_orig = torch.clamp(torch.randn(k_input_shape, dtype=dtypes[dtype]), min=-1, max=1).cuda()
     v_orig = torch.clamp(torch.randn(v_input_shape, dtype=dtypes[dtype]), min=-1, max=1).cuda()
@@ -468,11 +504,21 @@ def run_dpa_with_cp(
             x.view(*x.shape[:seq_dim], -1, *x.shape[(seq_dim + 2) :]) for x in [q_, k_, v_, dout_]
         ]
     elif qkv_format == "thd":
-        seq_idx_q = tex.thd_get_partitioned_indices(
-            cu_seqlens_q_padded, q_.shape[0], world_size, rank
+        seq_idx_q = get_thd_partitioned_indices(
+            cu_seqlens_q_padded,
+            q_.shape[0],
+            world_size,
+            rank,
+            device=q_.device,
+            load_balancing_strategy=load_balancing_strategy,
         )
-        seq_idx_kv = tex.thd_get_partitioned_indices(
-            cu_seqlens_kv_padded, k_.shape[0], world_size, rank
+        seq_idx_kv = get_thd_partitioned_indices(
+            cu_seqlens_kv_padded,
+            k_.shape[0],
+            world_size,
+            rank,
+            device=k_.device,
+            load_balancing_strategy=load_balancing_strategy,
         )
         q_, dout_ = [x.index_select(0, seq_idx_q) for x in [q_, dout_]]
         k_, v_ = [x.index_select(0, seq_idx_kv) for x in [k_, v_]]
@@ -516,6 +562,7 @@ def run_dpa_with_cp(
         cp_comm_ranks,
         torch.cuda.Stream(),
         cp_comm_type,
+        load_balancing_strategy,
     )
     if config.softmax_type != "vanilla":
         core_attn.softmax_offset.grad.zero_()

tests/pytorch/attention/test_attention_with_cp.py

@@ -708,3 +708,48 @@ def test_cp_with_fused_attention(
         deterministic=_deterministic,
         log_level=pytest_logging_level,
     )
+
+
+@pytest.mark.skipif(get_cudnn_version() < (8, 9, 7), reason="cuDNN 8.9.7+ is required.")
+@pytest.mark.skipif(
+    get_device_compute_capability() < (9, 0), reason="FusedAttention THD requires sm90+."
+)
+def test_cp_with_fused_attention_no_load_balance(cp_pool):
+    """Check experimental single-chunk forward/backward."""
+    # cp_2_0 reaches the generic test's known deterministic THD backward OOM threshold on sm90.
+    if _deterministic and get_device_compute_capability() == (9, 0):
+        pytest.skip(
+            "Deterministic FusedAttention backward with THD format OOMs on sm90"
+            " for large bHSS configs (known cuDNN issue)."
+        )
+    _submit(
+        cp_pool(2),
+        dtype="bf16",
+        model="cp_2_0",
+        qkv_format="thd",
+        kernel_backend="FusedAttention",
+        cp_comm_type="all_gather",
+        load_balancing_strategy="NO_LOAD_BALANCE",
+        deterministic=_deterministic,
+        log_level=pytest_logging_level,
+    )
+
+
+@pytest.mark.skipif(
+    get_device_compute_capability() != (9, 0) or not FlashAttentionUtils.v3_is_installed,
+    reason="FlashAttention 3 requires sm90 and an installed FA3 package.",
+)
+def test_cp_with_flash_attention_3_no_load_balance(cp_pool):
+    """Check the supported unpadded FlashAttention 3 path."""
+    _submit(
+        cp_pool(2),
+        dtype="bf16",
+        model="cp_2_0",
+        qkv_format="thd",
+        kernel_backend="FlashAttention",
+        cp_comm_type="all_gather",
+        fa_pad_between_seqs=False,
+        load_balancing_strategy="NO_LOAD_BALANCE",
+        deterministic=_deterministic,
+        log_level=pytest_logging_level,
+    )

tests/pytorch/attention/test_cp_utils.py

@@ -7,8 +7,13 @@
 import itertools
 import torch
 import unittest
+from transformer_engine.pytorch import CPLoadBalancingStrategy
 from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
+    get_no_load_balance_thd_causal_metadata,
     get_batch_on_this_cp_rank,
+    get_thd_partitioned_indices,
+    restore_thd_gathered_kv,
+    unrestore_thd_gathered_kv,
     pad_thd_sequences_for_cp,
     generate_positional_ids_for_cp,
 )
@@ -19,6 +24,49 @@
     tex = None
 
 
+class TestTHDPartitioning(unittest.TestCase):
+    def test_default_partition_rejects_cpu_metadata(self):
+        with self.assertRaisesRegex(AssertionError, "requires CUDA cu_seqlens"):
+            get_thd_partitioned_indices(torch.tensor([0, 8]), 8, 2, 0)
+
+    def test_no_load_balance_metadata_handles_document_padding_boundary(self):
+        cu_seqlens = torch.tensor([0, 6, 10], dtype=torch.int32)
+        cu_seqlens_padded = torch.tensor([0, 8, 12], dtype=torch.int32)
+
+        q_cu, q_cu_padded, kv_cu = get_no_load_balance_thd_causal_metadata(
+            cu_seqlens,
+            cu_seqlens_padded,
+            total_tokens=12,
+            cp_size=2,
+            cp_rank=1,
+        )
+
+        self.assertEqual(len(q_cu), 1)
+        self.assertTrue(torch.equal(q_cu[0], torch.tensor([0, 0, 4], dtype=torch.int32)))
+        self.assertTrue(torch.equal(q_cu_padded[0], torch.tensor([0, 2, 6], dtype=torch.int32)))
+        self.assertTrue(torch.equal(kv_cu[0], torch.tensor([0, 6, 10], dtype=torch.int32)))
+
+    def test_no_load_balance_restore_uses_captured_mode(self):
+        tokens = torch.arange(8)
+        cu_seqlens_padded = torch.tensor([0, 8], dtype=torch.int32)
+
+        restored = restore_thd_gathered_kv(
+            tokens,
+            cu_seqlens_padded,
+            2,
+            CPLoadBalancingStrategy.NO_LOAD_BALANCE,
+        )
+        unrestored = unrestore_thd_gathered_kv(
+            tokens,
+            cu_seqlens_padded,
+            2,
+            CPLoadBalancingStrategy.NO_LOAD_BALANCE,
+        )
+
+        self.assertIs(restored, tokens)
+        self.assertIs(unrestored, tokens)
+
+
 class TestSequencePadding(unittest.TestCase):
     def test_padding_with_custom_padding_values_sequences_shorter_than_divisibility_factor(
         self,
@@ -584,6 +632,46 @@ def test_cp_rank_slicing_simple_case(self):
         self.assertTrue(torch.equal(labels_r1, expected_labels_r1))
         self.assertTrue(torch.equal(pos_ids_r1, expected_pos_ids_r1))
 
+    @unittest.skipUnless(torch.cuda.is_available() and tex is not None, "CUDA extension required")
+    def test_cp_rank_slicing_dual_chunk_swap_on_cuda(self):
+        """CUDA inputs use the native DualChunkSwap partition indices."""
+        input_ids = torch.arange(16, device="cuda").unsqueeze(0)
+        labels = input_ids + 100
+        position_ids = torch.arange(16, device="cuda")
+        cu_seqlens = torch.tensor([0, 8, 16])
+
+        self._mock_distributed_env(cp_size=2, cp_rank=0)
+        input_ids_rank, labels_rank, position_ids_rank = get_batch_on_this_cp_rank(
+            cu_seqlens, input_ids, labels, position_ids
+        )
+
+        expected_indices = torch.tensor([0, 1, 6, 7, 8, 9, 14, 15], device="cuda")
+        self.assertTrue(torch.equal(input_ids_rank, input_ids.index_select(1, expected_indices)))
+        self.assertTrue(torch.equal(labels_rank, labels.index_select(1, expected_indices)))
+        self.assertTrue(
+            torch.equal(position_ids_rank, position_ids.index_select(0, expected_indices))
+        )
+
+    def test_cp_rank_slicing_no_load_balance_on_cpu(self):
+        """The experimental policy assigns one contiguous CPU chunk per rank."""
+        input_ids = torch.arange(12).unsqueeze(0)
+        labels = input_ids + 100
+        position_ids = torch.arange(12)
+        cu_seqlens = torch.tensor([0, 5, 12])
+
+        self._mock_distributed_env(cp_size=4, cp_rank=2)
+        input_ids_rank, labels_rank, position_ids_rank = get_batch_on_this_cp_rank(
+            cu_seqlens,
+            input_ids,
+            labels,
+            position_ids,
+            load_balancing_strategy=CPLoadBalancingStrategy.NO_LOAD_BALANCE,
+        )
+
+        self.assertTrue(torch.equal(input_ids_rank, torch.tensor([[6, 7, 8]])))
+        self.assertTrue(torch.equal(labels_rank, torch.tensor([[106, 107, 108]])))
+        self.assertTrue(torch.equal(position_ids_rank, torch.tensor([6, 7, 8])))
+
     def test_cp_rank_slicing_multiple_sequences(self):
         """Test CP rank slicing with multiple sequences."""
         # Setup: Two sequences of length 8 each, CP size = 2

transformer_engine/pytorch/__init__.py

@@ -17,7 +17,7 @@
 
 load_framework_extension("torch")
 from transformer_engine.pytorch import constants
-from transformer_engine.pytorch.constants import DType
+from transformer_engine.pytorch.constants import CPLoadBalancingStrategy, DType
 from transformer_engine.pytorch.module import LayerNormLinear
 from transformer_engine.pytorch.module import Linear
 from transformer_engine.pytorch.module import LayerNormMLP

transformer_engine/pytorch/attention/dot_product_attention/backends.py

@@ -37,6 +37,7 @@
 )
 from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor
 from transformer_engine.pytorch.constants import (
+    CPLoadBalancingStrategy,
     QKVLayouts,
     dist_group_type,
 )
@@ -900,6 +901,7 @@ def forward(
         num_splits: Optional[int] = 1,
         cu_seqlens_q_padded: Optional[torch.Tensor] = None,
         cu_seqlens_kv_padded: Optional[torch.Tensor] = None,
+        load_balancing_strategy=CPLoadBalancingStrategy.DUAL_CHUNK_SWAP,
     ) -> torch.Tensor:
         """flash-attn fprop"""
 
@@ -1134,6 +1136,7 @@ def forward(
                     pad_between_seqs=pad_between_seqs,
                     use_flash_attn_3=use_flash_attn_3,
                     fp8_output=fp8_output,
+                    load_balancing_strategy=load_balancing_strategy,
                 )
         else:
             if is_cpu_offload_enabled():
@@ -2116,6 +2119,7 @@ def forward(
         packed_qkv: Optional[torch.Tensor] = None,
         packed_kv: Optional[torch.Tensor] = None,
         bf16_backward: bool = False,
+        load_balancing_strategy=CPLoadBalancingStrategy.DUAL_CHUNK_SWAP,
     ) -> torch.Tensor:
         """fused attention fprop"""
         assert (
@@ -2282,6 +2286,7 @@ def forward(
                     fp8_output=fp8_output,
                     layer_number=self.layer_number,
                     return_max_logit=self.return_max_logit,
+                    load_balancing_strategy=load_balancing_strategy,
                 )
         elif score_mod is not None:
             output = FusedAttentionWithScoreModFunc.apply(