Add native transport for dynamic context parallelism

build_tools/pytorch.py

@@ -38,6 +38,13 @@ def setup_pytorch_extension(
 
     # Source files
     sources = all_files_in_dir(Path(csrc_source_files), name_extension="cpp")
+    build_native_cp_transport = bool(int(os.getenv("NVTE_WITH_NCCL_DEVICE_CP", "0")))
+    if build_native_cp_transport:
+        sources.extend(
+            path
+            for path in all_files_in_dir(Path(csrc_source_files), name_extension="cu")
+            if path.name == "cp_native_transport.cu"
+        )
 
     # Header files
     include_dirs = get_cuda_include_dirs()
@@ -87,16 +94,37 @@ def setup_pytorch_extension(
         libraries.append("nvshmem_host")
         cxx_flags.append("-DNVTE_ENABLE_NVSHMEM")
 
+    extra_compile_args = {"cxx": cxx_flags}
+    if build_native_cp_transport:
+        nvcc_flags = ["-O3", "-std=c++17"]
+        nccl_home = os.getenv("NCCL_HOME")
+        nccl_include_dir = os.getenv("NVTE_NCCL_INCLUDE_DIR")
+        nccl_library_dir = os.getenv("NVTE_NCCL_LIBRARY_DIR")
+        if nccl_home:
+            nccl_home = Path(nccl_home)
+            nccl_include_dir = nccl_include_dir or str(nccl_home / "include")
+            nccl_library_dir = nccl_library_dir or str(nccl_home / "lib")
+        if nccl_include_dir:
+            include_dirs.append(Path(nccl_include_dir))
+        if nccl_library_dir:
+            library_dirs.append(Path(nccl_library_dir))
+        libraries.append("nccl")
+        cxx_flags.append("-DNVTE_WITH_NCCL_DEVICE_CP")
+        nvcc_flags.append("-DNVTE_WITH_NCCL_DEVICE_CP")
+        extra_compile_args["nvcc"] = nvcc_flags
+
     # Construct PyTorch CUDA extension
     sources = [str(path) for path in sources]
     include_dirs = [str(path) for path in include_dirs]
-    from torch.utils.cpp_extension import CppExtension
+    from torch.utils.cpp_extension import CppExtension, CUDAExtension
+
+    extension_cls = CUDAExtension if build_native_cp_transport else CppExtension
 
-    return CppExtension(
+    return extension_cls(
         name="transformer_engine_torch",
         sources=[str(src) for src in sources],
         include_dirs=[str(inc) for inc in include_dirs],
-        extra_compile_args={"cxx": cxx_flags},
+        extra_compile_args=extra_compile_args,
         libraries=[str(lib) for lib in libraries],
         library_dirs=[str(lib_dir) for lib_dir in library_dirs],
     )

tests/pytorch/attention/run_attention_with_cp.py

@@ -5,15 +5,38 @@
 import os
 import sys
 import logging
+import copy
 from contextlib import nullcontext
 import torch
 import torch.distributed as dist
 from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
     get_cu_seqlens_on_cp_rank,
 )
+from transformer_engine.pytorch.attention.native_cp_transport import (
+    destroy_native_cp_transport,
+    initialize_native_cp_transport,
+    set_native_cp_parent_group,
+)
 from transformer_engine.pytorch.attention.dot_product_attention.utils import combine_and_quantize
 import transformer_engine_torch as tex
 from test_attention_with_cp import model_configs_flash_attn, model_configs_fused_attn
+
+
+class _LogicalCPGroup:
+    """Minimal topology descriptor used by native-transport tests."""
+
+    def __init__(self, ranks, rank):
+        self.ranks = tuple(ranks)
+        self.cp_size = len(self.ranks)
+        self.cp_rank = self.ranks.index(rank)
+
+    def size(self):
+        return self.cp_size
+
+    def rank(self):
+        return self.cp_rank
+
+
 from transformer_engine.pytorch import (
     autocast,
     DotProductAttention,
@@ -180,6 +203,9 @@ def run_dpa_with_cp(
     scaling_mode="delayed",
     f16_O="False",
     is_training="True",
+    native_cp_transport="False",
+    logical_cp_ring="False",
+    max_seqlen=None,
     log_level=logging.WARNING,
 ):
     """Test DotProductAttention module with context parallelism"""
@@ -202,6 +228,14 @@ def run_dpa_with_cp(
     if kernel_backend == "FusedAttention":
         os.environ["NVTE_FUSED_ATTN"] = "1"
         config = model_configs_fused_attn[model]
+    config = copy.deepcopy(config)
+    native_cp_transport = native_cp_transport == "True"
+    logical_cp_ring = logical_cp_ring == "True"
+    if logical_cp_ring and not native_cp_transport:
+        raise ValueError("logical_cp_ring requires native_cp_transport=True")
+    if max_seqlen is not None:
+        config.max_seqlen_q = int(max_seqlen)
+        config.max_seqlen_kv = int(max_seqlen)
     assert config.attn_mask_type in [
         "causal",
         "no_mask",
@@ -229,6 +263,14 @@ def run_dpa_with_cp(
     cp_comm_ranks = range(world_size)
     assert rank in cp_comm_ranks
     cp_comm_group = dist.new_group(cp_comm_ranks, backend="nccl")
+    cp_group = cp_comm_group
+    cp_rank = rank
+    cp_global_ranks = tuple(cp_comm_ranks)
+    if logical_cp_ring:
+        offsets = [0, *range(1, world_size, 2), *reversed(range(2, world_size, 2))]
+        cp_global_ranks = tuple(cp_comm_ranks[offset] for offset in offsets)
+        cp_group = _LogicalCPGroup(cp_global_ranks, rank)
+        cp_rank = cp_group.rank()
     if cp_comm_type == "a2a+p2p":
         assert world_size % 2 == 0, (
             "{cp_comm_type=} requires world_size % 2 = 0 as it assumes the a2a level has cp_size"
@@ -390,17 +432,17 @@ def run_dpa_with_cp(
             )
             for x in [q_, k_, v_, dout_]
         ]
-        seq_idx = torch.tensor([rank, 2 * world_size - rank - 1], device=q_.device)
+        seq_idx = torch.tensor([cp_rank, 2 * world_size - cp_rank - 1], device=q_.device)
         q_, k_, v_, dout_ = [x.index_select(seq_dim, seq_idx) for x in [q_, k_, v_, dout_]]
         q_, k_, v_, dout_ = [
             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
+            cu_seqlens_q_padded, q_.shape[0], world_size, cp_rank
         )
         seq_idx_kv = tex.thd_get_partitioned_indices(
-            cu_seqlens_kv_padded, k_.shape[0], world_size, rank
+            cu_seqlens_kv_padded, k_.shape[0], world_size, cp_rank
         )
         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_]]
@@ -438,10 +480,27 @@ def run_dpa_with_cp(
             bias_ = bias_.index_select(seq_q_dim, bias_seq_idx)
             bias_ = bias_.view(*shape_before_seq, -1, seq_kv_size)
             bias_.requires_grad = True
+
+    if native_cp_transport:
+        kv_bytes = (k_.numel() + v_.numel()) * k_.element_size()
+        pair_bytes = 2 * kv_bytes
+        gin_env_names = (
+            "NCCL_GIN_NCONTEXTS",
+            "NCCL_GIN_SIGNAL_POOL_SIZE",
+            "NCCL_GIN_COUNTER_POOL_SIZE",
+        )
+        gin_env_before = {name: os.environ.get(name) for name in gin_env_names}
+        initialize_native_cp_transport(
+            cp_comm_group,
+            ((pair_bytes + 255) // 256) * 256 + pair_bytes,
+        )
+        assert {name: os.environ.get(name) for name in gin_env_names} == gin_env_before
+        if logical_cp_ring:
+            set_native_cp_parent_group(cp_group, cp_comm_group)
     # set up environment
     core_attn.set_context_parallel_group(
-        cp_comm_sub_groups if cp_comm_type == "a2a+p2p" else cp_comm_group,
-        cp_comm_ranks,
+        cp_comm_sub_groups if cp_comm_type == "a2a+p2p" else cp_group,
+        cp_global_ranks,
         torch.cuda.Stream(),
         cp_comm_type,
     )
@@ -562,7 +621,7 @@ def run_dpa_with_cp(
             dq_, dk_, dv_, out_ = [dq_, dk_, dv_, out_]
             cu_seqlens_q_padded = cu_seqlens_q_padded // world_size
             cu_seqlens_q = get_cu_seqlens_on_cp_rank(
-                cu_seqlens_q, cu_seqlens_q_padded, world_size, rank, True, True
+                cu_seqlens_q, cu_seqlens_q_padded, world_size, cp_rank, True, True
             )
             cu_pads_q = cu_seqlens_q_padded - cu_seqlens_q
             num_pads_q = cu_pads_q[1:] - cu_pads_q[:-1]
@@ -582,7 +641,7 @@ def run_dpa_with_cp(
                     )
             cu_seqlens_kv_padded = cu_seqlens_kv_padded // world_size
             cu_seqlens_kv = get_cu_seqlens_on_cp_rank(
-                cu_seqlens_kv, cu_seqlens_kv_padded, world_size, rank, True, True
+                cu_seqlens_kv, cu_seqlens_kv_padded, world_size, cp_rank, True, True
             )
             cu_pads_kv = cu_seqlens_kv_padded - cu_seqlens_kv
             num_pads_kv = cu_pads_kv[1:] - cu_pads_kv[:-1]
@@ -733,6 +792,9 @@ def run_dpa_with_cp(
                 )
             logging.info(f"[Rank {rank}] CP vs no-CP: {names[i]} matches")
 
+    if native_cp_transport:
+        destroy_native_cp_transport(cp_comm_group)
+
     # destroy distribution group
     dist.destroy_process_group()
 

tests/pytorch/attention/test_native_cp_transport.py

@@ -0,0 +1,113 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""CPU-only checks for native CP initialization policy."""
+
+import importlib.util
+import os
+from pathlib import Path
+import sys
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+
+@pytest.fixture
+def native_module(monkeypatch):
+    """Load the transport without importing GPU-only TE modules."""
+    extension = Mock(cp_native_transport_create=Mock(return_value=(1, None)))
+    monkeypatch.setitem(sys.modules, "transformer_engine_torch", extension)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
+    monkeypatch.setattr(torch.distributed, "barrier", lambda **kwargs: None)
+    monkeypatch.setattr(torch.distributed, "get_process_group_ranks", lambda group: [0])
+    path = (
+        Path(__file__).resolve().parents[3]
+        / "transformer_engine/pytorch/attention/native_cp_transport.py"
+    )
+    spec = importlib.util.spec_from_file_location("native_cp_transport_policy_test", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module, extension
+
+
+@pytest.mark.parametrize("configured_value", [None, "512"])
+def test_native_cp_preserves_nccl_environment(monkeypatch, native_module, configured_value):
+    """Creating a transport must not limit GIN resources for other communicators."""
+    gin_env_names = (
+        "NCCL_GIN_NCONTEXTS",
+        "NCCL_GIN_SIGNAL_POOL_SIZE",
+        "NCCL_GIN_COUNTER_POOL_SIZE",
+    )
+    for name in gin_env_names:
+        if configured_value is None:
+            monkeypatch.delenv(name, raising=False)
+        else:
+            monkeypatch.setenv(name, configured_value)
+    module, extension = native_module
+    parent = Mock()
+    parent._get_backend.return_value._comm_ptr.return_value = 123
+
+    module.NativeCPTransport(parent, 256)
+
+    extension.cp_native_transport_create.assert_called_once_with(123, 256)
+    assert {name: os.environ.get(name) for name in gin_env_names} == {
+        name: configured_value for name in gin_env_names
+    }
+
+
+@pytest.mark.parametrize("dtype", [torch.int64, torch.bool, torch.float32])
+@pytest.mark.parametrize("peers", [(7, 11), (None, 11), (7, None), (None, None)])
+def test_native_cp_halo_stages_noncontiguous_tensors(native_module, dtype, peers):
+    """Halo staging reuses the arena and preserves missing-receive boundary fills."""
+    module, extension = native_module
+    transport = module.NativeCPTransport.__new__(module.NativeCPTransport)
+    transport.handle = 123
+    transport.arena = torch.empty(512, dtype=torch.uint8)
+    transport._parent_rank = {7: 1, 11: 2}
+    source = torch.arange(6).reshape(3, 2).to(dtype).t()
+    output = torch.zeros(3, 2, dtype=dtype).t()
+    expected = torch.ones_like(output)
+
+    def exchange(handle, send, recv, send_peer, recv_peer, channel):
+        assert handle == 123 and channel == 3
+        assert send_peer == (-1 if peers[0] is None else 1)
+        assert recv_peer == (-1 if peers[1] is None else 2)
+        assert send.is_contiguous() and recv.is_contiguous()
+        assert send.untyped_storage().data_ptr() == transport.arena.data_ptr()
+        assert recv.untyped_storage().data_ptr() == transport.arena.data_ptr()
+        if peers[0] is not None:
+            torch.testing.assert_close(send, source, rtol=0, atol=0)
+        if peers[1] is not None:
+            recv.copy_(expected)
+        return channel
+
+    extension.cp_native_transport_send_recv.side_effect = exchange
+    transport.exchange(source, peers[0], output, peers[1])
+    torch.testing.assert_close(
+        output, expected if peers[1] is not None else torch.zeros_like(output), rtol=0, atol=0
+    )
+    if peers == (None, None):
+        extension.cp_native_transport_send_recv.assert_not_called()
+        extension.cp_native_transport_wait.assert_not_called()
+    else:
+        extension.cp_native_transport_send_recv.assert_called_once()
+        extension.cp_native_transport_wait.assert_called_once_with(123, 3)
+
+
+def test_native_cp_halo_rejects_incompatible_buffers(native_module):
+    """Bad halo metadata must fail before entering the communication kernel."""
+    module, extension = native_module
+    transport = module.NativeCPTransport.__new__(module.NativeCPTransport)
+    transport.handle = 123
+    transport.arena = torch.empty(512, dtype=torch.uint8)
+    transport._parent_rank = {7: 1}
+    with pytest.raises(ValueError, match="matching shapes and dtypes"):
+        transport.exchange(torch.zeros(2), 7, torch.zeros(3), 7)
+    with pytest.raises(ValueError, match="matching shapes and dtypes"):
+        transport.exchange(torch.zeros(2), 7, torch.zeros(2, dtype=torch.int64), 7)
+    with pytest.raises(RuntimeError, match="arena needs"):
+        transport.exchange(torch.zeros(512), 7, torch.zeros(512), 7)
+    transport.exchange(torch.empty(0), 7, torch.empty(0), 7)
+    extension.cp_native_transport_send_recv.assert_not_called()

transformer_engine/pytorch/attention/dot_product_attention/backends.py

@@ -48,7 +48,10 @@
     META_QKV,
 )
 from transformer_engine.pytorch.quantization import get_fp8_torch_dtype, FP8GlobalStateManager
-from transformer_engine.pytorch.distributed import get_distributed_world_size
+from transformer_engine.pytorch.distributed import (
+    get_distributed_world_size,
+    is_logical_process_group,
+)
 from transformer_engine.pytorch.jit import no_torch_dynamo
 from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
     attn_forward_func_with_cp,
@@ -757,7 +760,7 @@ def forward(
         ), f"FlashAttention does not support qkv_layout = {qkv_layout}!"
 
         cp_size = 1
-        if isinstance(cp_group, dist_group_type):
+        if isinstance(cp_group, dist_group_type) or is_logical_process_group(cp_group):
             cp_size = get_distributed_world_size(cp_group)
         elif isinstance(cp_group, list):
             for group in cp_group:
@@ -1828,7 +1831,7 @@ def forward(
         ), f"FusedAttention does not support qkv_layout = {qkv_layout}!"
 
         cp_size = 1
-        if isinstance(cp_group, dist_group_type):
+        if isinstance(cp_group, dist_group_type) or is_logical_process_group(cp_group):
             cp_size = get_distributed_world_size(cp_group)
         elif isinstance(cp_group, list):
             for group in cp_group: