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,20 @@ 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
+        initialize_native_cp_transport(
+            cp_comm_group,
+            ((pair_bytes + 255) // 256) * 256 + pair_bytes,
+        )
+        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 +614,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 +634,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 +785,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()
 

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: