perf(comm): LL128 protocol for decode_cp_a2a_lse_reduce, on any torch symmetric-memory backend

benchmarks/bench_dcp_lse_reduce.py

@@ -3,6 +3,9 @@
 Example:
   torchrun --standalone --nproc-per-node=4 \
     benchmarks/bench_dcp_lse_reduce.py --label fused
+
+The workspace uses the process's default torch symmetric-memory backend unless
+--symm-mem-backend selects one.
 """
 
 import argparse
@@ -13,6 +16,7 @@
 
 import torch
 import torch.distributed as dist
+import torch.distributed._symmetric_memory as symm_mem
 
 from flashinfer.comm import (
     decode_cp_a2a_lse_reduce,
@@ -238,6 +242,9 @@ def nccl_baseline_call() -> torch.Tensor:
         "local_heads": local_heads,
         "head_dim": head_dim,
         "dtype": str(dtype).removeprefix("torch."),
+        "symm_mem_backend": symm_mem.get_backend(
+            torch.device("cuda", torch.cuda.current_device())
+        ),
         "eager_max_rank_median_us": _max_rank_median(eager_us, world_size),
         "nccl_baseline_max_rank_median_us": _max_rank_median(
             nccl_baseline_us, world_size
@@ -259,13 +266,21 @@ def main() -> None:
     parser.add_argument("--samples", type=int, default=30)
     parser.add_argument("--launches-per-sample", type=int, default=50)
     parser.add_argument("--graph-launches-per-replay", type=int, default=50)
+    parser.add_argument(
+        "--symm-mem-backend",
+        choices=["CUDA", "NCCL", "NVSHMEM"],
+        help="torch symmetric-memory backend for the workspace (default: the process's)",
+    )
     args = parser.parse_args()
     if args.graph_launches_per_replay <= 0:
         parser.error("--graph-launches-per-replay must be greater than zero")
 
     local_rank = int(os.environ["LOCAL_RANK"])
     device = torch.device("cuda", local_rank)
     torch.cuda.set_device(device)
+    # torch only accepts set_backend() for a backend other than its default.
+    if args.symm_mem_backend not in (None, symm_mem.get_backend(device)):
+        symm_mem.set_backend(args.symm_mem_backend)
     # Eagerly create the NCCL communicator before the symmetric-memory
     # rendezvous. This registers the group's host communicator for this
     # concrete CUDA device on released PyTorch builds.

docs/api/comm.rst

@@ -608,12 +608,13 @@ DCP All-to-All (Context-Parallel Attention Reduction)
     decode_cp_a2a_init_workspace
     decode_cp_a2a_alltoall
 
-NCCL LSA DCP All-to-All + LSE Reduce
--------------------------------------
+Fused DCP All-to-All + LSE Reduce
+---------------------------------
 
-The fused path uses PyTorch NCCL symmetric memory and requires every context-
-parallel rank to be in one load/store-accessible NVLink domain. It does not use
-the MNNVL workspace accepted by ``decode_cp_a2a_alltoall``.
+The fused path writes into every rank's workspace through PyTorch symmetric
+memory, on whichever backend the process uses, and requires every
+context-parallel rank to share one NVLink domain. It does not use the MNNVL
+workspace accepted by ``decode_cp_a2a_alltoall``.
 
 .. autosummary::
     :toctree: ../generated

flashinfer/comm/dcp_lse_reduce.py

@@ -16,9 +16,11 @@
 
 Public names follow https://github.com/flashinfer-ai/flashinfer/issues/4575.
 
-The implementation is self-contained in FlashInfer and uses NCCL's device-side
-LSA APIs over a tensor allocated by torch's NCCL symmetric-memory backend. It
-does not call ``libnccl_longseq`` and does not use the Helix/MNNVL A2A kernel.
+The implementation is self-contained in FlashInfer. Each rank writes straight
+into its peers' workspaces, a tensor allocated by torch symmetric memory on
+whichever backend the process uses (CUDA, NCCL, or NVSHMEM); it never selects a
+backend itself. It does not call ``libnccl_longseq`` and does not use the
+Helix/MNNVL A2A kernel.
 """
 
 import functools
@@ -83,10 +85,8 @@ def _get_workspace_state(workspace: torch.Tensor) -> _WorkspaceState:
     return state
 
 
-def _dcp_lse_reduce_payload_offset(cp_size: int) -> int:
-    # 16-byte epoch metadata plus two uint32 readiness arrays, rounded so the
-    # vectorized payload remains 16-byte aligned.
-    return (16 + 2 * cp_size * 4 + 15) // 16 * 16
+# Epoch metadata, padded so every 128-byte payload line is aligned.
+_PAYLOAD_OFFSET = 128
 
 
 @functools.cache
@@ -102,7 +102,7 @@ def decode_cp_a2a_lse_reduce_workspace_size(
     head_dim: int,
     dtype: torch.dtype,
 ) -> int:
-    """Return the required NCCL symmetric workspace size in bytes."""
+    """Return the required symmetric-memory workspace size in bytes."""
     if min(max_tokens, local_heads, cp_size, head_dim) <= 0:
         raise ValueError("workspace geometry must be positive")
     if cp_size > 64:
@@ -112,10 +112,14 @@ def decode_cp_a2a_lse_reduce_workspace_size(
     itemsize = torch.empty((), dtype=dtype).element_size()
     if head_dim * itemsize % 16 != 0:
         raise ValueError("head_dim rows must be 16-byte aligned")
-    payload_offset = _dcp_lse_reduce_payload_offset(cp_size)
-    return payload_offset + 2 * cp_size * max_tokens * local_heads * (
-        head_dim * itemsize + 4
-    )
+    # Each row is sent as 128-byte lines of 15 eight-byte data words plus a flag
+    # word, and its LSE as one 16-byte line per group of four data lines.
+    lines = (head_dim * itemsize // 8 + 14) // 15
+    groups = (lines + 3) // 4
+    if groups > 32:
+        raise ValueError("head_dim rows must not exceed 15360 bytes")
+    row_bytes = lines * 128 + groups * 16
+    return _PAYLOAD_OFFSET + 2 * cp_size * max_tokens * local_heads * row_bytes
 
 
 @flashinfer_api
@@ -127,14 +131,19 @@ def decode_cp_a2a_lse_reduce_create_workspace(
     dtype: torch.dtype,
     group: Any,
 ) -> torch.Tensor:
-    """Create and rendezvous the fused op's NCCL symmetric workspace.
-
-    All ranks in ``group`` must call this function collectively. The group must
-    fit in one NCCL load/store-accessible (LSA) NVLink domain; multi-node groups
-    spanning LSA domains are not supported. Allocate one workspace per group
-    and reuse it for every invocation and CUDA graph replay. A workspace may
-    only be used from one ordered CUDA stream; allocate a workspace per
-    concurrent stream.
+    """Create and rendezvous the fused op's symmetric-memory workspace.
+
+    All ranks in ``group`` must call this function collectively. The workspace
+    is allocated on the process's current torch symmetric-memory backend
+    (``torch.distributed._symmetric_memory.get_backend``; CUDA unless the
+    application selected another one), which this function never changes. Every
+    rank's workspace must be mapped for load/store over NVLink: the group must
+    fit in one node or one multi-node NVLink domain. With the NCCL backend,
+    ``group`` must be an NCCL group whose communicator already exists (e.g. after
+    one collective on it), and cuMem must not be disabled. Allocate one
+    workspace per group and reuse it for every invocation and CUDA graph
+    replay. A workspace may only be used from one ordered CUDA stream; allocate
+    a workspace per concurrent stream.
 
     Parameters
     ----------
@@ -154,31 +163,36 @@ def decode_cp_a2a_lse_reduce_create_workspace(
     Returns
     -------
     torch.Tensor
-        A rendezvoused NCCL symmetric-memory ``uint8`` tensor. Allocate it once
+        A rendezvoused symmetric-memory ``uint8`` tensor. Allocate it once
         before CUDA graph capture and reuse it.
     """
-    # Build/load on every rank before entering either NCCL collective below.
-    # The JIT cache lock may otherwise leave one rank creating the devcomm
-    # while another rank is still compiling.
+    # Build/load on every rank before the collective rendezvous below, so that
+    # no rank waits there on another that is still compiling.
     get_dcp_lse_reduce_module()
     size_bytes = decode_cp_a2a_lse_reduce_workspace_size(
         max_tokens, local_heads, cp_size, head_dim, dtype
     )
     group_name = group if isinstance(group, str) else group.group_name
-    symm_mem.set_backend("NCCL")
     # PyTorch's NCCL symmetric-memory communicator registry is keyed by the
     # concrete CUDA device. Do not use the unindexed ``"cuda"`` device here:
     # it would rendezvous through a separate registry entry from the process
     # group's ``cuda:<local_rank>`` communicator.
     device = torch.device("cuda", torch.cuda.current_device())
     workspace = symm_mem.empty(size_bytes, dtype=torch.uint8, device=device)
-    # Initialize the local epoch and readiness words. The payload is fully
-    # overwritten before every read; clearing it could race a peer's first put.
-    workspace[: _dcp_lse_reduce_payload_offset(cp_size)].zero_()
+    # Data lines are ready when their flag word matches the (nonzero) per-call
+    # flag, so the whole workspace, not only the epoch, must start zeroed.
+    workspace.zero_()
     # Initialization must complete before rendezvous; afterwards any peer may
-    # enter the first fused kernel and publish a remote readiness value.
-    torch.cuda.current_stream().synchronize()
+    # enter the first fused kernel and write lines into this workspace.
+    torch.cuda.current_stream(device).synchronize()
     handle = symm_mem.rendezvous(workspace, group)
+    peer_ptrs = handle.buffer_ptrs
+    if len(peer_ptrs) != handle.world_size or not all(peer_ptrs):
+        raise RuntimeError(
+            "decode_cp_a2a_lse_reduce needs every rank's workspace mapped for "
+            f"load/store, but the {symm_mem.get_backend(device)} symmetric-memory "
+            "backend could not map all of them; the group must share one NVLink domain"
+        )
     state = _WorkspaceState(handle=handle, group_name=group_name, device=device)
     data_ptr = workspace.data_ptr()
     _workspace_keepalive[data_ptr] = state
@@ -195,28 +209,34 @@ def decode_cp_a2a_lse_reduce(
     cp_size: int,
     lse_mode: Literal["base2", "basee"] = "base2",
 ) -> torch.Tensor:
-    """Fuse an NCCL LSA DCP A2A exchange with the LSE-weighted reduce.
+    """Fuse the DCP A2A exchange with the LSE-weighted reduce.
 
     Send, receive synchronization, and reduction execute in one cooperative
-    CUDA kernel.
+    CUDA kernel. Data moves in 128-byte lines that carry their own readiness
+    flag (NCCL's LL128 protocol), so no fence or grid-wide barrier is needed.
 
     This is a collective operation. Every rank must invoke it in the same order
     and the same number of times, including CUDA graph replays. All ranks must
-    belong to one NCCL LSA/NVLink domain.
+    share one NVLink domain.
 
     Parameters
     ----------
     partial_o : torch.Tensor
         ``[batch, heads, cp_size, head_dim]`` CUDA tensor (fp16 or bf16), or
         more generally ``[..., cp_size, head_dim]``.
         ``partial_o[..., peer, :]`` is the slice destined for that CP rank.
+        A 4-D input may be any strided view whose last dimension is contiguous,
+        e.g. an attention output ``[batch, cp_size * heads, head_dim]`` viewed as
+        ``out.unflatten(1, (cp_size, heads)).permute(0, 2, 1, 3)``, so no packing
+        copy is needed. Inputs of any other dimensionality must be contiguous.
     partial_lse : torch.Tensor
         ``[batch, heads, cp_size]`` CUDA float32 tensor, or more generally
         ``[..., cp_size]``. Its leading dimensions must match ``partial_o``.
+        Like ``partial_o``, a 3-D input may be a strided view.
         ``heads`` is simply the number of heads present in the input; it may
         be local or total because this operation does not shard the head axis.
     workspace : torch.Tensor
-        Rendezvoused NCCL symmetric-memory tensor from
+        Rendezvoused symmetric-memory tensor from
         :func:`decode_cp_a2a_lse_reduce_create_workspace`. Reuse it only from
         one ordered CUDA stream.
     cp_rank : int

flashinfer/jit/comm.py

@@ -425,16 +425,16 @@ def gen_dcp_alltoall_module() -> JitSpec:
 
 
 def gen_dcp_lse_reduce_module() -> JitSpec:
-    """Build the NCCL-symmetric-memory DCP A2A + LSE-reduce op."""
+    """Build the DCP A2A + LSE-reduce op over torch symmetric memory."""
     from torch.utils.cpp_extension import include_paths, library_paths
 
     nvcc_flags = current_compilation_context.get_nvcc_flags_list(
         supported_major_versions=[9, 10, 11, 12]
     )
-    # This kernel uses cooperative groups and NCCL's device API, neither of
-    # which requires Hopper's architecture-specific ``90a`` ISA. Compile a
-    # generic SM90 image so it also loads on H100 installations that expose
-    # compute capability 9.0 without the ``a`` feature set.
+    # This kernel needs only a cooperative launch, not Hopper's
+    # architecture-specific ``90a`` ISA. Compile a generic SM90 image so it also
+    # loads on H100 installations that expose compute capability 9.0 without the
+    # ``a`` feature set.
     nvcc_flags = [
         flag.replace("compute_90a,code=sm_90a", "compute_90,code=sm_90")
         for flag in nvcc_flags
@@ -450,67 +450,18 @@ def gen_dcp_lse_reduce_module() -> JitSpec:
         pathlib.Path(path) for path in cuda_include_paths
     ]
     extra_ldflags = [f"-L{path}" for path in cuda_library_paths]
-    nccl_ldflag = "-lnccl"
-
-    nccl_home_path: pathlib.Path | None = None
-    nccl_home = os.environ.get("NCCL_HOME")
-    if nccl_home:
-        nccl_home_path = pathlib.Path(nccl_home)
-    else:
-        # PyTorch CUDA wheels install NCCL as a Python package, but
-        # cpp_extension does not add that package's headers or libraries to
-        # extension builds. Discover the wheel location directly so CI and
-        # ordinary pip environments do not need to export NCCL_HOME.
-        try:
-            import nvidia.nccl  # type: ignore[import-not-found]
-
-            nccl_home_path = pathlib.Path(nvidia.nccl.__path__[0])
-        except (ImportError, AttributeError, IndexError):
-            # A system NCCL installation may already be on the compiler and
-            # linker search paths; retain the -lnccl fallback below.
-            pass
-
-    if nccl_home_path is not None:
-        nccl_lib_candidates = [
-            nccl_home_path / "lib",
-            nccl_home_path / "lib64",
-            nccl_home_path,
-        ]
-        nccl_lib_dir = next(
-            (
-                path
-                for path in nccl_lib_candidates
-                if (path / "libnccl.so").exists() or any(path.glob("libnccl.so.*"))
-            ),
-            nccl_home_path / "lib",
-        )
-        extra_includes.extend(
-            [
-                nccl_home_path / "include",
-                nccl_home_path / "include" / "nccl_device",
-            ]
-        )
-        extra_ldflags.insert(0, f"-L{nccl_lib_dir}")
-        # PyTorch's pip NCCL wheels include only the versioned shared object,
-        # whereas system installations also provide libnccl.so for -lnccl.
-        if not (nccl_lib_dir / "libnccl.so").exists():
-            versioned_libraries = sorted(nccl_lib_dir.glob("libnccl.so.*"))
-            if versioned_libraries:
-                nccl_ldflag = str(versioned_libraries[-1])
-
     extra_ldflags += [
         "-ltorch",
         "-ltorch_cpu",
         "-ltorch_cuda",
         "-lc10",
         "-lc10_cuda",
-        nccl_ldflag,
     ]
 
     return gen_jit_spec(
         "dcp_lse_reduce",
         [jit_env.FLASHINFER_CSRC_DIR / "dcp_lse_reduce.cu"],
         extra_include_paths=extra_includes,
-        extra_cuda_cflags=["-std=c++20", "-DUSE_NCCL"] + nvcc_flags,
+        extra_cuda_cflags=["-std=c++20"] + nvcc_flags,
         extra_ldflags=extra_ldflags,
     )

flashinfer/trace/templates/comm.py

@@ -432,7 +432,7 @@ def pcie_ipc_reduce_scatter_trace_dispatch(**kwargs):
 )
 
 
-# ── Fused NCCL-LSA DCP all-to-all + LSE reduce ───────────────────────────────
+# ── Fused DCP all-to-all + LSE reduce ────────────────────────────────────────
 
 
 @torch.no_grad()
@@ -494,7 +494,7 @@ def _decode_cp_a2a_lse_reduce_init(
     op_type="comm",
     name_prefix="decode_cp_a2a_lse_reduce",
     description=(
-        "Fused context-parallel NCCL-LSA all-to-all and LSE-weighted "
+        "Fused context-parallel all-to-all and LSE-weighted "
         "attention-output reduction. The trace reference models the local "
         "LSE merge; multi-rank exchange correctness is exercised by tests/comm."
     ),