diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py
index 7c6cdefd15..b620d7e4fe 100644
--- a/tests/pytorch/attention/run_attention_with_cp.py
+++ b/tests/pytorch/attention/run_attention_with_cp.py
@@ -11,9 +11,9 @@
 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,
@@ -21,6 +21,7 @@
 )
 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_()
diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py
index d7eb16b862..5ac4dd4578 100644
--- a/tests/pytorch/attention/test_attention_with_cp.py
+++ b/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,
+    )
diff --git a/tests/pytorch/attention/test_cp_utils.py b/tests/pytorch/attention/test_cp_utils.py
index c3a423cef5..81ee4781d3 100644
--- a/tests/pytorch/attention/test_cp_utils.py
+++ b/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
diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py
index 2b1803bfb2..40816c0797 100644
--- a/transformer_engine/pytorch/__init__.py
+++ b/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
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index 8a219a6a4d..a6a1e56be3 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/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(
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
index ea89ca97eb..c69ce1adca 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
@@ -25,7 +25,10 @@
 from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage
 from transformer_engine.pytorch.jit import jit_fuser
 from transformer_engine.pytorch.graph import is_graph_capturing
-from transformer_engine.pytorch.constants import dist_group_type
+from transformer_engine.pytorch.constants import (
+    CPLoadBalancingStrategy,
+    dist_group_type,
+)
 from transformer_engine.pytorch.distributed import (
     get_distributed_world_size,
     get_distributed_rank,
@@ -267,6 +270,111 @@ def get_seq_chunk_ids_for_reordering_after_attn(cp_size, device):
     return _seq_chunk_ids_cache_for_reordering_after_attn[(cp_size, device)]
 
 
+def _get_thd_partition_cu_seqlens(cu_seqlens_padded, device=None):
+    """Return physical boundaries used by the THD partition CUDA kernels."""
+    target_device = torch.device(device if device is not None else cu_seqlens_padded.device)
+    target_dtype = torch.int32 if target_device.type == "cuda" else cu_seqlens_padded.dtype
+    return cu_seqlens_padded.to(device=target_device, dtype=target_dtype)
+
+
+def get_thd_partitioned_indices(
+    cu_seqlens_padded,
+    total_tokens,
+    cp_size,
+    cp_rank,
+    device=None,
+    load_balancing_strategy=CPLoadBalancingStrategy.DUAL_CHUNK_SWAP,
+):
+    """Return THD token indices using the selected CP partition contract."""
+    assert isinstance(
+        load_balancing_strategy, CPLoadBalancingStrategy
+    ), f"Expected {CPLoadBalancingStrategy.__name__}, got {type(load_balancing_strategy).__name__}."
+    if load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE:
+        validate_no_load_balance_thd_metadata(
+            cu_seqlens_padded,
+            cu_seqlens_padded,
+            total_tokens,
+            cp_size,
+        )
+        target_device = torch.device(device if device is not None else cu_seqlens_padded.device)
+        target_dtype = torch.int32 if target_device.type == "cuda" else cu_seqlens_padded.dtype
+        chunk_size = total_tokens // cp_size
+        return torch.arange(
+            cp_rank * chunk_size,
+            (cp_rank + 1) * chunk_size,
+            dtype=target_dtype,
+            device=target_device,
+        )
+    cu_seqlens_padded = _get_thd_partition_cu_seqlens(cu_seqlens_padded, device)
+    assert cu_seqlens_padded.is_cuda, (
+        "Per-document THD partitioning requires CUDA cu_seqlens; pass device='cuda' "
+        "when the source metadata is on CPU."
+    )
+    if cu_seqlens_padded.dtype != torch.int32:
+        cu_seqlens_padded = cu_seqlens_padded.to(torch.int32)
+    return tex.thd_get_partitioned_indices(cu_seqlens_padded, total_tokens, cp_size, cp_rank)
+
+
+def validate_no_load_balance_thd_metadata(
+    cu_seqlens,
+    cu_seqlens_padded,
+    total_tokens,
+    cp_size,
+):
+    """Validate no-load-balance THD metadata while producing rank-local inputs."""
+    assert cu_seqlens.shape == cu_seqlens_padded.shape
+    assert total_tokens % cp_size == 0
+    assert cu_seqlens[0] == 0 and cu_seqlens_padded[0] == 0
+    assert cu_seqlens_padded[-1] == total_tokens
+    assert torch.all(cu_seqlens[1:] >= cu_seqlens[:-1])
+    assert torch.all(cu_seqlens_padded[1:] >= cu_seqlens_padded[:-1])
+    actual_seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
+    padded_seqlens = cu_seqlens_padded[1:] - cu_seqlens_padded[:-1]
+    assert torch.all(actual_seqlens <= padded_seqlens)
+
+
+def get_no_load_balance_thd_causal_metadata(
+    cu_seqlens,
+    cu_seqlens_padded,
+    total_tokens,
+    cp_size,
+    cp_rank,
+):
+    """Build one-step THD metadata for no-load-balance CP partitioning.
+
+    The complete physical token buffer is the sharding unit. ``cu_seqlens``
+    remains the logical document boundary, so each global chunk is represented
+    as its intersection with every document.
+    """
+    assert cp_size > 0 and 0 <= cp_rank < cp_size
+    assert cu_seqlens.shape == cu_seqlens_padded.shape
+    assert total_tokens % cp_size == 0
+
+    chunk_size = total_tokens // cp_size
+    chunk_start = cp_rank * chunk_size
+    chunk_end = chunk_start + chunk_size
+    actual_seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
+    doc_starts = cu_seqlens_padded[:-1]
+    valid_doc_ends = doc_starts + actual_seqlens
+
+    fragment_starts = torch.clamp(doc_starts, min=chunk_start, max=chunk_end)
+    fragment_ends = torch.clamp(valid_doc_ends, min=chunk_start, max=chunk_end)
+    fragment_seqlens = torch.clamp(fragment_ends - fragment_starts, min=0)
+
+    q_cu_seqlens = torch.zeros_like(cu_seqlens)
+    q_cu_seqlens[1:] = fragment_seqlens.cumsum(0)
+    q_cu_seqlens_padded = (
+        torch.clamp(cu_seqlens_padded, min=chunk_start, max=chunk_end) - chunk_start
+    )
+
+    visible_kv_seqlens = torch.clamp(chunk_end - doc_starts, min=0)
+    visible_kv_seqlens = torch.minimum(visible_kv_seqlens, actual_seqlens)
+    kv_cu_seqlens = torch.zeros_like(cu_seqlens)
+    kv_cu_seqlens[1:] = visible_kv_seqlens.cumsum(0)
+
+    return [q_cu_seqlens], [q_cu_seqlens_padded], [kv_cu_seqlens]
+
+
 @jit_fuser
 def reorder_seq_chunks_for_a2a_before_attn(x, chunk_ids_for_a2a, seq_dim, cp_size):
     """Reorder sequence chunk for A2A communication before attention compute."""
@@ -384,6 +492,24 @@ def thd_cp_rank_order_to_sequence_order(x, cu_seqlens, cp_size, seq_dim=0):
     return tex.thd_cp_rank_order_to_sequence_order(x, cu_seqlens, cp_size, x.shape[seq_dim])
 
 
+def restore_thd_gathered_kv(x, cu_seqlens_padded, cp_size, load_balancing_strategy):
+    """Restore gathered THD tokens using the strategy captured by attention forward."""
+    if load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE:
+        # Rank r owns physical chunk r, so rank-major all-gather is already in sequence order.
+        return x
+    cu_seqlens_padded = _get_thd_partition_cu_seqlens(cu_seqlens_padded, x.device)
+    return thd_cp_rank_order_to_sequence_order(x, cu_seqlens_padded, cp_size)
+
+
+def unrestore_thd_gathered_kv(x, cu_seqlens_padded, cp_size, load_balancing_strategy):
+    """Arrange THD tokens for reduce-scatter using the captured attention strategy."""
+    if load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE:
+        # Physical sequence order is also the rank-major reduce-scatter order for this policy.
+        return x
+    cu_seqlens_padded = _get_thd_partition_cu_seqlens(cu_seqlens_padded, x.device)
+    return thd_sequence_order_to_cp_rank_order(x, cu_seqlens_padded, cp_size)
+
+
 def flash_attn_a2a_communicate(
     a2a_inputs: Union[torch.Tensor, List[torch.Tensor]],
     chunk_ids_for_a2a: torch.Tensor,
@@ -3056,6 +3182,7 @@ def forward(
         fp8_meta,
         quantizers,
         fp8_output,
+        load_balancing_strategy,
     ):
         # pylint: disable=missing-function-docstring
         nvtx_range_push("transformer_engine.AttnFuncWithCPAndKVAllGather.forward")
@@ -3094,10 +3221,11 @@ def forward(
             f" >= 2.3. Found {use_fused_attention=}, {use_flash_attn_3=}, "
             f"and {fa_utils.v2_3_plus=}."
         )
-        assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, (
-            "cp_comm_type='all_gather' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q ="
-            f" {q.shape[seq_dim_qkv]}, seq_len_kv = {k.shape[seq_dim_qkv]}."
-        )
+        if load_balancing_strategy is CPLoadBalancingStrategy.DUAL_CHUNK_SWAP:
+            assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, (
+                "cp_comm_type='all_gather' requires seq_len % 2 == 0 for Q, K, V. Found "
+                f"seq_len_q = {q.shape[seq_dim_qkv]}, seq_len_kv = {k.shape[seq_dim_qkv]}."
+            )
 
         flash_attn_fwd = None
         if not use_fused_attention:
@@ -3142,14 +3270,21 @@ def forward(
                 q.shape[seq_dim] % 2 == 0 and k.shape[seq_dim] % 2 == 0
             ), "Sequence length per GPU needs to be divisible by 2!"
 
-        # Divide by 2*cp_size to get per-chunk values
-        max_seqlen_q = max_seqlen_q // (2 * cp_size)
-        max_seqlen_kv = max_seqlen_kv // (2 * cp_size)
+        # Per-document DCS divides every sequence into 2*CP chunks. No-load-balance
+        # instead bounds Q by one global chunk and keeps full-document KV bounds.
+        if load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE:
+            max_seqlen_q = min(max_seqlen_q, q.shape[0])
+        else:
+            max_seqlen_q = max_seqlen_q // (2 * cp_size)
+            max_seqlen_kv = max_seqlen_kv // (2 * cp_size)
         if use_fused_attention and qkv_format != "thd":
             cu_seqlens_q = cu_seqlens_q // (2 * cp_size)
-        if qkv_format == "thd":
+        if (
+            qkv_format == "thd"
+            and load_balancing_strategy is CPLoadBalancingStrategy.DUAL_CHUNK_SWAP
+        ):
             cu_seqlens_q_padded = cu_seqlens_q_padded // (2 * cp_size)
-        else:
+        elif qkv_format != "thd":
             cu_seqlens_q_padded = None
         if use_fused_attention and attn_mask_type == "causal":
             attn_mask_type = attn_mask_type + "_bottom_right"
@@ -3217,10 +3352,12 @@ def forward(
 
         if qkv_format == "thd":
             # [cp*t, h, d] -> reorder to sequence order -> [t_full, h, d]
-            # The padded cu_seqlens are global sequence offsets. Reorder uses them to
-            # derive per-sequence chunk boundaries.
-            k_ag = thd_cp_rank_order_to_sequence_order(k_ag, cu_seqlens_kv_padded, cp_size)
-            v_ag = thd_cp_rank_order_to_sequence_order(v_ag, cu_seqlens_kv_padded, cp_size)
+            k_ag = restore_thd_gathered_kv(
+                k_ag, cu_seqlens_kv_padded, cp_size, load_balancing_strategy
+            )
+            v_ag = restore_thd_gathered_kv(
+                v_ag, cu_seqlens_kv_padded, cp_size, load_balancing_strategy
+            )
         else:
             # [cp, s, b, h, d] -> [cp*2, s//2, b, h, d]
             k_ag = k_ag.view(2 * cp_size, k.shape[0] // 2, *k.shape[1:])
@@ -3231,12 +3368,9 @@ def forward(
             # [cp*2, s//2, b, h, d] -> [cp*s, b, h, d]
             k_ag = k_ag.view(-1, *k.shape[1:])
             v_ag = v_ag.view(-1, *v.shape[1:])
-        # cp_stream is used for step 1 of the per-step loop and must wait until
-        # k_ag/v_ag preparation finishes on the current stream — otherwise step 1
-        # races against AG/reorder writes. Manifests at high cp_size where reorder
-        # is large enough to outlast cp_stream's launch (e.g. bucket128k @ cp=8).
-        cp_stream.wait_stream(torch.cuda.current_stream())
-
+            # Non-THD cp_stream inputs are ready after K/V reorder, so wait here to
+            # preserve overlap with output initialization below.
+            cp_stream.wait_stream(torch.cuda.current_stream())
         # THD all_gather only reaches this path for f16/bf16 attention today.
         # q: [b, 2, s//2, h, d] or [2, s//2, b, h, d]
         # k: [s, b, h, d]
@@ -3251,7 +3385,11 @@ def forward(
         # create two streams to resolve wave quantization issue of Flash Attn in each step
         flash_attn_streams = [torch.cuda.current_stream(), cp_stream]
         # prepare per-step tensors
-        local_seq_chunk_ids = [rank, 2 * cp_size - rank - 1]
+        local_seq_chunk_ids = (
+            [rank]
+            if load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE
+            else [rank, 2 * cp_size - rank - 1]
+        )
         kv_seq_range_per_step = [None, None]
         window_size_per_step = [None, None]
         cu_seqlens_kv_per_step = [None, None]
@@ -3264,8 +3402,30 @@ def forward(
         max_logit_per_step = [None, None]
         max_logit = None
 
+        # Initialize before the conditional so static analysis can prove they are
+        # assigned before the backend-specific loop below.
+        thd_cu_seqlens_q_per_step = [None, None]
+        thd_cu_seqlens_q_padded_per_step = [None, None]
+        thd_cu_seqlens_kv_per_step = [None, None]
+
         # Pre-compute THD-specific per-step cu_seqlens
-        if qkv_format == "thd":
+        if (
+            qkv_format == "thd"
+            and load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE
+        ):
+            total_tokens_q = q.shape[0] * cp_size
+            (
+                thd_cu_seqlens_q_per_step,
+                thd_cu_seqlens_q_padded_per_step,
+                thd_cu_seqlens_kv_per_step,
+            ) = get_no_load_balance_thd_causal_metadata(
+                cu_seqlens_q_original,
+                cu_seqlens_q_padded,
+                total_tokens_q,
+                cp_size,
+                rank,
+            )
+        elif qkv_format == "thd":
             # Rank-level padded offsets (2 chunks per sequence on this rank)
             cu_seqlens_q_padded_rank = cu_seqlens_q_padded * 2
 
@@ -3336,6 +3496,11 @@ def forward(
                 thd_cu_seqlens_kv_per_step[0][1:] = visible_actual[0].cumsum(0)
                 thd_cu_seqlens_kv_per_step[1][1:] = visible_actual[1].cumsum(0)
 
+        if qkv_format == "thd":
+            # Delay the THD wait so one dependency covers both restored K/V and the
+            # per-step metadata produced above on the current stream.
+            cp_stream.wait_stream(torch.cuda.current_stream())
+
         for i in range(len(local_seq_chunk_ids) + 1):
             if i < len(local_seq_chunk_ids):
                 # FA3 uses internal per-call workspace. Consecutive AG per-step
@@ -3395,15 +3560,19 @@ def forward(
                         q_part = q
                         k_part = k_ag
                         v_part = v_ag
-                        kv_range, window_size_per_step[i] = get_kv_seq_info_after_all_gather(
-                            local_seq_chunk_ids[i],
-                            cp_size,
-                            max_seqlen_q,
-                            max_seqlen_kv,
-                            window_size,
-                            causal,
-                        )
-                        max_seqlen_kv_ = kv_range[1]
+                        if load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE:
+                            window_size_per_step[i] = (-1, 0)
+                            max_seqlen_kv_ = max_seqlen_kv
+                        else:
+                            kv_range, window_size_per_step[i] = get_kv_seq_info_after_all_gather(
+                                local_seq_chunk_ids[i],
+                                cp_size,
+                                max_seqlen_q,
+                                max_seqlen_kv,
+                                window_size,
+                                causal,
+                            )
+                            max_seqlen_kv_ = kv_range[1]
                         cu_seqlens_kv_per_step[i] = thd_cu_seqlens_kv_per_step[i]
                     if use_fused_attention:
                         # Set per-step parameters for THD vs bshd/sbhd
@@ -3648,6 +3817,7 @@ def forward(
         ctx.use_flash_attn_3 = use_flash_attn_3
         ctx.pad_between_seqs = pad_between_seqs
         ctx.window_size = window_size
+        ctx.load_balancing_strategy = load_balancing_strategy
         if qkv_format == "thd":
             ctx.max_seqlen_kv = max_seqlen_kv
             ctx.cu_seqlens_kv_padded = cu_seqlens_kv_padded
@@ -3792,9 +3962,12 @@ def backward(ctx, dout, *_args):
             cu_seqlens_kv_padded = ctx.cu_seqlens_kv_padded
             thd_cu_seqlens_q_per_step = ctx.thd_cu_seqlens_q_per_step
             # [cp*t, h, d] -> reorder to sequence order
-            # Use padded cu_seqlens (divisible by 2*cp_size) for correct reorder
-            k_ag = thd_cp_rank_order_to_sequence_order(k_ag, cu_seqlens_kv_padded, cp_size)
-            v_ag = thd_cp_rank_order_to_sequence_order(v_ag, cu_seqlens_kv_padded, cp_size)
+            k_ag = restore_thd_gathered_kv(
+                k_ag, cu_seqlens_kv_padded, cp_size, ctx.load_balancing_strategy
+            )
+            v_ag = restore_thd_gathered_kv(
+                v_ag, cu_seqlens_kv_padded, cp_size, ctx.load_balancing_strategy
+            )
 
             thd_cu_seqlens_q_padded_per_step = ctx.thd_cu_seqlens_q_padded_per_step
         else:
@@ -3842,7 +4015,11 @@ def backward(ctx, dout, *_args):
                 if fa_utils.v2_6_0_plus:
                     fa_backward_kwargs["softcap"] = 0.0
 
-        local_seq_chunk_ids = [rank, 2 * cp_size - rank - 1]
+        local_seq_chunk_ids = (
+            [rank]
+            if ctx.load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE
+            else [rank, 2 * cp_size - rank - 1]
+        )
         for i in range(len(local_seq_chunk_ids) + 1):
             if i < len(local_seq_chunk_ids):
                 # FA3 uses internal per-call workspace. Consecutive AG per-step
@@ -3857,15 +4034,18 @@ def backward(ctx, dout, *_args):
                         q_part = q
                         k_part = k_ag
                         v_part = v_ag
-                        kv_range, _ = get_kv_seq_info_after_all_gather(
-                            local_seq_chunk_ids[i],
-                            cp_size,
-                            ctx.max_seqlen_q,
-                            ctx.max_seqlen_kv,
-                            ctx.window_size,
-                            "causal" in ctx.attn_mask_type,
-                        )
-                        max_seqlen_kv = kv_range[1]
+                        if ctx.load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE:
+                            max_seqlen_kv = ctx.max_seqlen_kv
+                        else:
+                            kv_range, _ = get_kv_seq_info_after_all_gather(
+                                local_seq_chunk_ids[i],
+                                cp_size,
+                                ctx.max_seqlen_q,
+                                ctx.max_seqlen_kv,
+                                ctx.window_size,
+                                "causal" in ctx.attn_mask_type,
+                            )
+                            max_seqlen_kv = kv_range[1]
                         out_part = out
                         dout_part = dout
                     else:
@@ -4101,9 +4281,12 @@ def backward(ctx, dout, *_args):
         if ctx.qkv_format == "thd":
             # Reorder dK/dV from sequence order back to dual-chunk CP rank order,
             # then reduce-scatter across CP ranks.
-            # Use padded cu_seqlens for correct slice boundaries.
-            dk = thd_sequence_order_to_cp_rank_order(dk, cu_seqlens_kv_padded, cp_size)
-            dv = thd_sequence_order_to_cp_rank_order(dv, cu_seqlens_kv_padded, cp_size)
+            dk = unrestore_thd_gathered_kv(
+                dk, cu_seqlens_kv_padded, cp_size, ctx.load_balancing_strategy
+            )
+            dv = unrestore_thd_gathered_kv(
+                dv, cu_seqlens_kv_padded, cp_size, ctx.load_balancing_strategy
+            )
             dk, _ = reduce_scatter_along_first_dim(dk, ctx.cp_group)
             dv, _ = reduce_scatter_along_first_dim(dv, ctx.cp_group)
             # dQ is already [t_rank, h, d], no reshape needed
@@ -4164,6 +4347,7 @@ def backward(ctx, dout, *_args):
             None,
             None,
             None,
+            None,
         )
 
 
@@ -4955,14 +5139,23 @@ def attn_forward_func_with_cp(
     fp8_output=False,
     layer_number=1,
     return_max_logit=False,
+    load_balancing_strategy=CPLoadBalancingStrategy.DUAL_CHUNK_SWAP,
 ) -> torch.Tensor:
     """
     Attention implementation with context parallelism (CP). CP partitions tensors along the sequence
     dimension, and by reducing the memory and computational pressure on each GPU, it enables long-context
-    LLMs in a distributed fashion. Transformer Engine's PyTorch CP implementation currently utilizes
-    the DualChunkSwap strategy to ensure load balancing across CP ranks. It is applied to all `attn_mask_type`s
-    and all `qkv_format`s, and it requires sequence lengths to be, or are padded to be, divisible by
-    (cp_size * 2). It also requires tokens to be re-ordered before entering this function.
+    LLMs in a distributed fashion. By default, Transformer Engine's PyTorch CP
+    implementation applies DualChunkSwap to each sequence independently. It requires
+    every sequence length to be, or be padded to be, divisible by (cp_size * 2), and
+    tokens must be re-ordered before entering this function.
+
+    Experimental strategy ``CPLoadBalancingStrategy.NO_LOAD_BALANCE`` instead
+    assigns one contiguous physical-buffer chunk to each rank and uses one attention
+    step per rank. Logical sequences remain isolated by ``cu_seqlens``. This strategy
+    requires THD, all-gather, causal self-attention without a sliding window, and
+    FusedAttention, or FlashAttention 3 with ``pad_between_seqs=False``. Input producers
+    must use the same strategy when partitioning inputs with
+    :func:`get_batch_on_this_cp_rank` or :func:`get_thd_partitioned_indices`.
 
     For qkv_format = {'bshd', 'sbhd'}, the token re-ordering is illustrated as below, for an example
     use case of s = 12, attn_mask_type = 'causal', and cp_size = 2. seq_pos indicates each token's position
@@ -5040,6 +5233,37 @@ def attn_forward_func_with_cp(
         "sbhd",
         "thd",
     ], f"Context parallelism does not support {qkv_format=}!"
+    assert isinstance(
+        load_balancing_strategy, CPLoadBalancingStrategy
+    ), f"Expected {CPLoadBalancingStrategy.__name__}, got {type(load_balancing_strategy).__name__}."
+    if load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE:
+        assert qkv_format == "thd", "No-load-balance CP partitioning requires qkv_format='thd'."
+        assert (
+            cp_comm_type == "all_gather"
+        ), "No-load-balance THD partitioning requires cp_comm_type='all_gather'."
+        assert (
+            use_fused_attention or use_flash_attn_3
+        ), "No-load-balance THD partitioning requires FusedAttention or FlashAttention 3."
+        assert not (
+            use_flash_attn_3 and pad_between_seqs
+        ), "No-load-balance THD partitioning with FlashAttention 3 does not support padding yet."
+        assert "causal" in attn_mask_type and window_size == (
+            -1,
+            0,
+        ), (
+            "No-load-balance THD partitioning requires causal attention without a sliding "
+            "window (window_size=(-1, 0))."
+        )
+        assert not fp8, "No-load-balance THD partitioning does not support FP8 yet."
+        assert (
+            not is_graph_capturing()
+        ), "No-load-balance THD partitioning does not support CUDA graph capture yet."
+        assert (
+            q.shape[0] == k.shape[0] == v.shape[0]
+        ), "No-load-balance THD partitioning requires equal local Q/K/V physical lengths."
+        assert cu_seqlens_q is cu_seqlens_kv and (
+            cu_seqlens_q_padded is cu_seqlens_kv_padded
+        ), "No-load-balance THD self-attention requires shared Q/KV sequence metadata tensors."
     assert (
         qkv_format != "sbhd" or use_fused_attention
     ), "Context parallelism does not support FlashAttention backend with qkv_format = 'sbhd'!"
@@ -5118,6 +5342,7 @@ def attn_forward_func_with_cp(
             fp8_meta,
             quantizers,
             fp8_output,
+            load_balancing_strategy,
         ]
         out = AttnFuncWithCPAndKVAllGather.apply(*args)
     elif cp_comm_type == "a2a":
@@ -5261,6 +5486,7 @@ def get_batch_on_this_cp_rank(
     position_ids_padded: torch.Tensor,
     cp_group: torch.distributed.ProcessGroup = None,
     qvk_format: str = "thd",
+    load_balancing_strategy=CPLoadBalancingStrategy.DUAL_CHUNK_SWAP,
 ):
     """Slice batch input along sequence dimension into multiple chunks for THD format.
 
@@ -5269,32 +5495,80 @@ def get_batch_on_this_cp_rank(
 
     Which are parallelized across GPUs in a context parallel group.
     This version works with variable-length sequences using cumulative sequence lengths.
+    The experimental ``CPLoadBalancingStrategy.NO_LOAD_BALANCE`` strategy assigns one
+    contiguous physical-buffer chunk per rank. By default, each padded sequence is chunked
+    independently.
     """
     if qvk_format not in ["thd", "bshd", "sbhd"]:
         raise ValueError(f"Unsupported qvk_format: {qvk_format}!")
+    assert isinstance(
+        load_balancing_strategy, CPLoadBalancingStrategy
+    ), f"Expected {CPLoadBalancingStrategy.__name__}, got {type(load_balancing_strategy).__name__}."
     if qvk_format == "thd":
         # Get context parallel size and rank
         cp_size = torch.distributed.get_world_size(group=cp_group)
         if cp_size > 1:
             cp_rank = torch.distributed.get_rank(group=cp_group)
+            seq_len_val = cu_seqlens_padded[-1].item()
+            rank_indices_by_device = {}
+
+            if load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE:
+
+                def build_rank_indices(device):
+                    return get_thd_partitioned_indices(
+                        cu_seqlens_padded,
+                        seq_len_val,
+                        cp_size,
+                        cp_rank,
+                        device,
+                        load_balancing_strategy,
+                    )
+
+            else:
+                total_slices = 2 * cp_size
+                slice_sizes = (cu_seqlens_padded[1:] - cu_seqlens_padded[:-1]) // total_slices
+
+                def build_rank_indices(device):
+                    if device.type == "cuda":
+                        return get_thd_partitioned_indices(
+                            cu_seqlens_padded,
+                            seq_len_val,
+                            cp_size,
+                            cp_rank,
+                            device,
+                            load_balancing_strategy,
+                        )
+                    # Preserve the CPU-capable per-document dataloader path.
+                    rank_slices = []
+                    for slice_size, seq_start in zip(slice_sizes, cu_seqlens_padded[:-1]):
+                        rank_slices.extend(
+                            [
+                                torch.arange(
+                                    seq_start + cp_rank * slice_size,
+                                    seq_start + (cp_rank + 1) * slice_size,
+                                    device=device,
+                                ),
+                                torch.arange(
+                                    seq_start + (total_slices - cp_rank - 1) * slice_size,
+                                    seq_start + (total_slices - cp_rank) * slice_size,
+                                    device=device,
+                                ),
+                            ]
+                        )
+                    return torch.cat(rank_slices)
 
-            # Calculate the chunk sizes for each sequence
-            total_slices_of_any_sequence = 2 * cp_size
-            slice_sizes = (
-                cu_seqlens_padded[1:] - cu_seqlens_padded[:-1]
-            ) // total_slices_of_any_sequence
+            def get_rank_indices(device):
+                """Build partition indices once for each input device."""
+                device = torch.device(device)
+                if device not in rank_indices_by_device:
+                    rank_indices_by_device[device] = build_rank_indices(device)
+                return rank_indices_by_device[device]
 
             # Process each tensor directly instead of using keys_to_change loop
             def process_tensor(val):
                 if val is None:
                     return val
                 # Determine which dimension is the sequence dimension
-                # Ensure cu_seqlens_padded[-1] is a Python int, not a 0-dim tensor
-                if isinstance(cu_seqlens_padded[-1], torch.Tensor):
-                    seq_len_val = cu_seqlens_padded[-1].item()
-                else:
-                    seq_len_val = cu_seqlens_padded[-1]
-
                 # Handle 1D tensors (like position_ids that don't have batch dimension)
                 if val.ndim == 1:
                     if val.shape[0] == seq_len_val:
@@ -5316,29 +5590,7 @@ def process_tensor(val):
                 else:
                     raise ValueError("Tensor must be at least 1D")
 
-                # On this particular rank, for each sequence, get two slices, one from the beginning
-                # and one from the end.
-                cp_rank_slices = []
-                for slice_size, seq_start in zip(slice_sizes, cu_seqlens_padded[:-1]):
-                    # 1st segment
-                    cp_rank_slices.append(
-                        torch.arange(
-                            seq_start + (cp_rank * slice_size),
-                            seq_start + ((cp_rank + 1) * slice_size),
-                            device=val.device,
-                        )
-                    )
-
-                    # 2nd segment
-                    cp_rank_slices.append(
-                        torch.arange(
-                            seq_start + ((total_slices_of_any_sequence - cp_rank - 1) * slice_size),
-                            seq_start + ((total_slices_of_any_sequence - cp_rank) * slice_size),
-                            device=val.device,
-                        )
-                    )
-
-                return val.index_select(current_seq_dim, torch.cat(cp_rank_slices))
+                return val.index_select(current_seq_dim, get_rank_indices(val.device))
 
             # Process each tensor directly
             input_ids_padded = process_tensor(input_ids_padded)
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index d5adbbcadf..1e5f5552d4 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -36,7 +36,13 @@
 from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage
 from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
 from transformer_engine.pytorch.export import is_in_onnx_export_mode
-from transformer_engine.pytorch.constants import AttnMaskTypes, AttnTypes, dist_group_type, DType
+from transformer_engine.pytorch.constants import (
+    AttnMaskTypes,
+    AttnTypes,
+    CPLoadBalancingStrategy,
+    DType,
+    dist_group_type,
+)
 from transformer_engine.pytorch.distributed import (
     get_distributed_world_size,
     checkpoint,
@@ -727,6 +733,7 @@ def __init__(
         self.cp_global_ranks = cp_global_ranks
         self.cp_stream = cp_stream
         self.cp_comm_type = cp_comm_type
+        self.load_balancing_strategy = CPLoadBalancingStrategy.DUAL_CHUNK_SWAP
 
         self.hidden_size_per_attention_head_k = (
             kv_channels if isinstance(kv_channels, int) else kv_channels[0]
@@ -879,6 +886,7 @@ def set_context_parallel_group(
         cp_global_ranks: List[int],
         cp_stream: torch.cuda.Stream,
         cp_comm_type: str = "p2p",
+        load_balancing_strategy=CPLoadBalancingStrategy.DUAL_CHUNK_SWAP,
     ) -> None:
         """
         Set the context parallel attributes for the given
@@ -908,11 +916,18 @@ def set_context_parallel_group(
                       - ``"a2a+p2p"``: hierarchical CP implementation. First applying a2a to QKV
                         across each CP sub-group (e.g., via NVLink), then exchanging KV with
                         p2p between sub-groups (e.g., via IBLink).
+        load_balancing_strategy : CPLoadBalancingStrategy
+                                  token partition strategy for context-parallel attention.
         """
+        assert isinstance(load_balancing_strategy, CPLoadBalancingStrategy), (
+            f"Expected {CPLoadBalancingStrategy.__name__}, "
+            f"got {type(load_balancing_strategy).__name__}."
+        )
         self.cp_group = cp_group
         self.cp_global_ranks = cp_global_ranks
         self.cp_stream = cp_stream
         self.cp_comm_type = cp_comm_type
+        self.load_balancing_strategy = load_balancing_strategy
 
     def init_fp8_metadata(self, num_gemms: int = 1) -> None:
         """
@@ -2210,6 +2225,7 @@ def forward(
                     cp_global_ranks=self.cp_global_ranks,
                     cp_stream=self.cp_stream,
                     cp_comm_type=self.cp_comm_type,
+                    load_balancing_strategy=self.load_balancing_strategy,
                     max_seqlen_q=max_seqlen_q,
                     max_seqlen_kv=max_seqlen_kv,
                     fp8=self.fp8 and self.fp8_meta["recipe"].fp8_dpa,
@@ -2266,6 +2282,7 @@ def forward(
                         cp_global_ranks=self.cp_global_ranks,
                         cp_stream=self.cp_stream,
                         cp_comm_type=self.cp_comm_type,
+                        load_balancing_strategy=self.load_balancing_strategy,
                         fp8=self.fp8 and self.fp8_meta["recipe"].fp8_dpa,
                         fp8_meta=self.fp8_meta,
                         quantizers=self.quantizers,
@@ -2300,6 +2317,7 @@ def forward(
                     cp_global_ranks=self.cp_global_ranks,
                     cp_stream=self.cp_stream,
                     cp_comm_type=self.cp_comm_type,
+                    load_balancing_strategy=self.load_balancing_strategy,
                     fp8=self.fp8 and self.fp8_meta["recipe"].fp8_dpa,
                     fp8_meta=self.fp8_meta,
                     quantizers=self.quantizers,
diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py
index f87365bf7c..ed68dc82f0 100644
--- a/transformer_engine/pytorch/attention/multi_head_attention.py
+++ b/transformer_engine/pytorch/attention/multi_head_attention.py
@@ -22,6 +22,7 @@
 from transformer_engine.pytorch.constants import (
     AttnTypes,
     AttnBiasTypes,
+    CPLoadBalancingStrategy,
     dist_group_type,
 )
 from transformer_engine.pytorch.distributed import (
@@ -659,6 +660,7 @@ def set_context_parallel_group(
         cp_global_ranks: List[int],
         cp_stream: torch.cuda.Stream,
         cp_comm_type: str = "p2p",
+        load_balancing_strategy=CPLoadBalancingStrategy.DUAL_CHUNK_SWAP,
     ) -> None:
         """
         Set the context parallel attributes for the given
@@ -688,6 +690,8 @@ def set_context_parallel_group(
                       - ``"a2a+p2p"``: hierarchical CP implementation. First applying a2a to QKV
                         across each CP sub-group (e.g., via NVLink), then exchanging KV with
                         p2p between sub-groups (e.g., via IBLink).
+        load_balancing_strategy : CPLoadBalancingStrategy
+                                  token partition strategy for context-parallel attention.
         """
         if isinstance(cp_group, dist_group_type):
             self.cp_size = get_distributed_world_size(cp_group)
@@ -706,12 +710,24 @@ def set_context_parallel_group(
             self.cp_size = cp_size_a2a * cp_size_p2p
             self.cp_rank = cp_size_a2a * cp_rank_p2p + cp_rank_a2a
 
+        # Preserve the legacy child-setter call unless an experimental strategy is requested.
+        load_balancing_kwargs = (
+            {}
+            if load_balancing_strategy is CPLoadBalancingStrategy.DUAL_CHUNK_SWAP
+            else {"load_balancing_strategy": load_balancing_strategy}
+        )
         # Deep iterate but skip self to avoid infinite recursion.
         for index, child in enumerate(self.modules()):
             if index == 0:
                 continue
             if hasattr(child, "set_context_parallel_group"):
-                child.set_context_parallel_group(cp_group, cp_global_ranks, cp_stream, cp_comm_type)
+                child.set_context_parallel_group(
+                    cp_group,
+                    cp_global_ranks,
+                    cp_stream,
+                    cp_comm_type,
+                    **load_balancing_kwargs,
+                )
 
     def forward(
         self,
diff --git a/transformer_engine/pytorch/constants.py b/transformer_engine/pytorch/constants.py
index 3a145bbb5b..54f7f8144f 100644
--- a/transformer_engine/pytorch/constants.py
+++ b/transformer_engine/pytorch/constants.py
@@ -11,6 +11,16 @@
 import transformer_engine_torch as tex
 
 
+class CPLoadBalancingStrategy(enum.Enum):
+    """Token partition strategy for context-parallel attention.
+
+    ``NO_LOAD_BALANCE`` is experimental.
+    """
+
+    DUAL_CHUNK_SWAP = "dual_chunk_swap"
+    NO_LOAD_BALANCE = "no_load_balance"
+
+
 class DType(enum.IntEnum):
     """Transformer Engine data types used to tag tensors passed to the
     Transformer Engine backend.
diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py
index d377e5f3b3..56028396e3 100644
--- a/transformer_engine/pytorch/transformer.py
+++ b/transformer_engine/pytorch/transformer.py
@@ -28,6 +28,7 @@
 )
 from transformer_engine.pytorch.constants import (
     AttnMaskTypes,
+    CPLoadBalancingStrategy,
     LayerTypes,
     dist_group_type,
 )
@@ -594,6 +595,7 @@ def set_context_parallel_group(
         cp_global_ranks: List[int],
         cp_stream: torch.cuda.Stream,
         cp_comm_type: str = "p2p",
+        load_balancing_strategy=CPLoadBalancingStrategy.DUAL_CHUNK_SWAP,
     ) -> None:
         r"""
         Set the context parallel attributes for the given
@@ -623,13 +625,28 @@ def set_context_parallel_group(
                       - ``"a2a+p2p"``: hierarchical CP implementation. First applying a2a to QKV
                         across each CP sub-group (e.g., via NVLink), then exchanging KV with
                         p2p between sub-groups (e.g., via IBLink).
+        load_balancing_strategy : CPLoadBalancingStrategy
+                                  token partition strategy for context-parallel attention.
+                                  ``NO_LOAD_BALANCE`` is experimental.
         """
+        # Preserve the legacy child-setter call unless an experimental strategy is requested.
+        load_balancing_kwargs = (
+            {}
+            if load_balancing_strategy is CPLoadBalancingStrategy.DUAL_CHUNK_SWAP
+            else {"load_balancing_strategy": load_balancing_strategy}
+        )
         # Deep iterate but skip self to avoid infinite recursion.
         for index, child in enumerate(self.modules()):
             if index == 0:
                 continue
             if hasattr(child, "set_context_parallel_group"):
-                child.set_context_parallel_group(cp_group, cp_global_ranks, cp_stream, cp_comm_type)
+                child.set_context_parallel_group(
+                    cp_group,
+                    cp_global_ranks,
+                    cp_stream,
+                    cp_comm_type,
+                    **load_balancing_kwargs,
+                )
 
     def forward(
         self,