[JAX] Add EP borrowed-comm path via XLA FFI collectives extension

build_tools/jax.py

@@ -5,6 +5,7 @@
 """JAX related extensions."""
 
 import os
+import warnings
 from pathlib import Path
 from packaging import version
 
@@ -101,6 +102,23 @@ def setup_jax_extension(
         ]
     )
 
+    # TODO(nccl-ep): temporary WAR -- do not upstream. Remove once jaxlib ships
+    # xla/ffi/api/collectives_c_api.h in its include dir.
+    #
+    # Optional extra include root for XLA FFI headers current jaxlib omits,
+    # needed by the EP borrowed-comm path. Point NVTE_JAX_XLA_FFI_EXTRA_INCLUDE
+    # at an XLA source checkout to build that path today.
+    xla_ffi_extra_include = os.getenv("NVTE_JAX_XLA_FFI_EXTRA_INCLUDE")
+    if xla_ffi_extra_include:
+        extra_root = Path(xla_ffi_extra_include)
+        include_dirs.append(extra_root)
+        if not (extra_root / "xla" / "ffi" / "api" / "collectives_c_api.h").is_file():
+            warnings.warn(
+                "NVTE_JAX_XLA_FFI_EXTRA_INCLUDE is set to "
+                f"'{xla_ffi_extra_include}' but xla/ffi/api/collectives_c_api.h was "
+                "not found there; the EP borrowed-comm path will not be built."
+            )
+
     # Compile flags
     cxx_flags = ["-O3"]
     if debug_build_enabled():

qa/L2_jax_distributed_unittest/test.sh

@@ -15,4 +15,6 @@ mkdir -p "$XML_LOG_DIR"
 XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_triton_gemm=false" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
 
 # NCCL EP multi-process suite. The launcher skips when fewer than 4 GPUs or no NVLink is detected.
+# Runs the borrowed-comm suite too (L2 only).
+export NVTE_JAX_UNITTEST_LEVEL="L2"
 TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh

tests/jax/test_multi_process_ep.py

@@ -25,6 +25,7 @@
 import re
 import sys
 import unittest
+from unittest import mock
 
 import jax
 import jax.experimental.multihost_utils as jmu
@@ -47,8 +48,13 @@
     ep_dispatch_fwd,
     ep_combine_fwd,
     get_ep_config,
+    is_ep_borrowed_comm_built,
+    use_nccl_comm_from_xla,
+)
+from transformer_engine.jax.version_utils import (
+    is_collective_stream_supported,
+    is_xla_ffi_collectives_supported,
 )
-from transformer_engine.jax.version_utils import is_collective_stream_supported
 
 
 # ── Test config ─────────────────────────────────────────────────────────────
@@ -107,11 +113,23 @@ def _local_device_sm():
 
 
 class TestEP(unittest.TestCase):
+    # Selects the EP comm path for this class. False forces the self-hosted NCCL
+    # comm; the TestEPBorrowedComm subclass flips it to exercise the borrowed path.
+    USE_BORROWED_COMM = False
+
     @classmethod
     def setUpClass(cls):
         sm = _local_device_sm()
         if sm is not None and sm < 90:
             raise unittest.SkipTest(f"NCCL EP requires SM>=90 (got SM{sm})")
+        if cls.USE_BORROWED_COMM and not (
+            is_ep_borrowed_comm_built() and is_xla_ffi_collectives_supported()
+        ):
+            raise unittest.SkipTest("EP borrowed-comm path needs a newer JAX/XLA build")
+        cls._prev_comm_env = os.environ.get("NVTE_JAX_EP_NCCL_COMM_FROM_XLA")
+        os.environ["NVTE_JAX_EP_NCCL_COMM_FROM_XLA"] = "1" if cls.USE_BORROWED_COMM else "0"
+        # Drop any communicator a prior class left so we bootstrap on a clean slate.
+        ep_finalize()
         cls.num_procs = jax.process_count()
         cls.rank = jax.process_index()
         cls.dp, cls.ep = _factor_dp_ep(cls.num_procs)
@@ -144,6 +162,15 @@ def setUpClass(cls):
         # alignment exercises dispatch_output_per_expert_alignment end-to-end.
         cls.hk = EpLayerConfig(top_k=TOP_K, dispatch_output_per_expert_alignment=16)
 
+    @classmethod
+    def tearDownClass(cls):
+        # Leave a clean slate for the next class and restore the env override.
+        ep_finalize()
+        if cls._prev_comm_env is None:
+            os.environ.pop("NVTE_JAX_EP_NCCL_COMM_FROM_XLA", None)
+        else:
+            os.environ["NVTE_JAX_EP_NCCL_COMM_FROM_XLA"] = cls._prev_comm_env
+
     # ── Bootstrap precondition ────────────────────────────────────────────
 
     def test_bootstrap_rejects_missing_ep_axis(self):
@@ -820,6 +847,37 @@ def bwd_only(eo, toks, idx, w, g):
                 self.assertEqual(hlo.count(op), 0, f"unexpected XLA {op} in bwd HLO:\n{hlo}")
 
 
+# ── Borrowed-comm path ───────────────────────────────────────────────────────
+
+
+class TestEPBorrowedComm(TestEP):
+    """Re-run EP primitives on the XLA borrowed-comm path.
+
+    Skipped entirely unless the build and installed JAX both provide the
+    collectives FFI extension. To keep L0/L1 fast, only a small smoke subset
+    (_SMOKE) runs by default; the full borrowed-path suite runs at L2
+    (NVTE_JAX_UNITTEST_LEVEL=L2).
+    """
+
+    USE_BORROWED_COMM = True
+
+    # Representative cases kept outside L2: one dispatch/combine round-trip (fwd)
+    # and its gradient (bwd). Every other inherited case runs only at L2.
+    _SMOKE = frozenset(
+        {
+            "test_primitive_dispatch_combine_identity_uniform",
+            "test_primitive_dispatch_combine_identity_bwd_uniform",
+        }
+    )
+
+    def setUp(self):
+        if (
+            os.environ.get("NVTE_JAX_UNITTEST_LEVEL", "L0") != "L2"
+            and self._testMethodName not in self._SMOKE
+        ):
+            self.skipTest("borrowed-comm full suite runs at L2 (NVTE_JAX_UNITTEST_LEVEL=L2)")
+
+
 # ── Drop-on-overflow ─────────────────────────────────────────────────────────
 
 
@@ -961,6 +1019,46 @@ def test_ep_tp_splits_domains(self):
         self.assertEqual(domains, {0: [0, 2, 4, 6], 1: [1, 3, 5, 7]})
 
 
+# ── Comm-path selection (single-process; no GPU needed) ──────────────────────
+
+
+class TestEpCommSelection(unittest.TestCase):
+    """use_nccl_comm_from_xla() build/version gating and NVTE_JAX_EP_NCCL_COMM_FROM_XLA override."""
+
+    @staticmethod
+    def _use(env, built, supported):
+        import transformer_engine.jax.cpp_extensions.ep as ep_mod
+
+        prev = os.environ.pop("NVTE_JAX_EP_NCCL_COMM_FROM_XLA", None)
+        if env is not None:
+            os.environ["NVTE_JAX_EP_NCCL_COMM_FROM_XLA"] = env
+        try:
+            with mock.patch.object(
+                ep_mod, "is_ep_borrowed_comm_built", return_value=built
+            ), mock.patch.object(
+                ep_mod, "is_xla_ffi_collectives_supported", return_value=supported
+            ):
+                return ep_mod.use_nccl_comm_from_xla()
+        finally:
+            os.environ.pop("NVTE_JAX_EP_NCCL_COMM_FROM_XLA", None)
+            if prev is not None:
+                os.environ["NVTE_JAX_EP_NCCL_COMM_FROM_XLA"] = prev
+
+    def test_auto_requires_build_and_version(self):
+        # Env unset: borrowed path only when both build and JAX support it.
+        self.assertTrue(self._use(None, built=True, supported=True))
+        self.assertFalse(self._use(None, built=True, supported=False))
+        self.assertFalse(self._use(None, built=False, supported=True))
+
+    def test_env_override_wins_over_version(self):
+        self.assertTrue(self._use("1", built=True, supported=False))
+        self.assertFalse(self._use("0", built=True, supported=True))
+
+    def test_force_on_without_build_raises(self):
+        with self.assertRaisesRegex(RuntimeError, "without the EP borrowed-comm path"):
+            self._use("1", built=False, supported=True)
+
+
 # ── Entry point ──────────────────────────────────────────────────────────────
 
 
@@ -981,7 +1079,7 @@ def test_ep_tp_splits_domains(self):
     )
 
     loader = unittest.TestLoader()
-    test_cases = (TestEP, TestEPOverflowDrop, TestEpDomainGrouping)
+    test_cases = (TestEP, TestEPBorrowedComm, TestEPOverflowDrop, TestEpDomainGrouping)
     target = os.environ.get("TARGET_TEST")
     if target:
         name = target.split(".")[-1]

transformer_engine/jax/cpp_extensions/ep.py

@@ -15,17 +15,19 @@
 """
 
 import functools
+import os
 from dataclasses import dataclass
 
 import jax
 import jax.numpy as jnp
+import numpy as np
 from jax import dtypes, ffi
 from jax.sharding import NamedSharding, PartitionSpec
 
 import transformer_engine_jax
 from .base import BasePrimitive, register_primitive
 from ..sharding import global_mesh_resource, get_mesh_axis_size
-from ..version_utils import is_collective_stream_supported
+from ..version_utils import is_collective_stream_supported, is_xla_ffi_collectives_supported
 
 
 def _on_collective_stream(func):
@@ -69,6 +71,35 @@ def wrapper(*args, **kwargs):
 # ── Module-level EP config ──────────────────────────────────────────────────
 
 
+@functools.lru_cache(maxsize=None)
+def is_ep_borrowed_comm_built() -> bool:
+    """True if transformer_engine_jax was compiled with the borrowed-comm FFI."""
+    try:
+        return "te_ep_bootstrap_borrowed_comm_ffi" in transformer_engine_jax.registrations()
+    except Exception:  # pylint: disable=broad-except
+        return False
+
+
+def use_nccl_comm_from_xla() -> bool:
+    """True when EP should borrow XLA's NCCL comm instead of self-hosting NCCL.
+
+    Auto-selected when both the build and the installed JAX support the XLA
+    collectives FFI extension. NVTE_JAX_EP_NCCL_COMM_FROM_XLA=1/0 is an internal
+    override for tests, not a supported user knob.
+    """
+    env = os.environ.get("NVTE_JAX_EP_NCCL_COMM_FROM_XLA")
+    if env is not None:
+        forced_on = env not in ("0", "", "false", "False")
+        if forced_on and not is_ep_borrowed_comm_built():
+            raise RuntimeError(
+                "NVTE_JAX_EP_NCCL_COMM_FROM_XLA is set but transformer_engine_jax was built "
+                "without the EP borrowed-comm path (XLA collectives FFI headers were "
+                "unavailable at build time). Unset it to use the self-hosted NCCL comm."
+            )
+        return forced_on
+    return is_ep_borrowed_comm_built() and is_xla_ffi_collectives_supported()
+
+
 @dataclass(frozen=True)
 class EpConfig:
     """Snapshot of the EP bootstrap config (see ep_bootstrap).
@@ -91,6 +122,39 @@ class EpConfig:
 _ep_config: EpConfig = None
 
 
+# Fixed sentinel keeps EP on its own private comm so it never aliases an XLA
+# collective over the same devices. Must stay in [0, 2**63 - 1].
+# 0x54454550 spells "TEEP".
+EP_COMMUNICATION_ID = 0x54454550
+
+
+def run_borrowed_comm_bootstrap(
+    mesh, replica_groups_flat, group_size, communication_id=EP_COMMUNICATION_ID
+):
+    """Initialize EPBackend on the borrowed XLA comm (one-shot, all devices)."""
+    try:
+        from jax import shard_map  # top-level since v0.8.0
+    except ImportError:  # older JAX
+        from jax.experimental.shard_map import shard_map
+
+    all_axes = tuple(mesh.axis_names)
+    spec = PartitionSpec(all_axes)
+    world = int(np.prod([mesh.shape[a] for a in all_axes]))
+    rg = np.asarray(replica_groups_flat, np.int64)
+    gs = np.int64(group_size)
+    cid = np.int64(communication_id)
+
+    def _body(x):
+        out_type = jax.ShapeDtypeStruct(x.shape, x.dtype)
+        return ffi.ffi_call("te_ep_bootstrap_borrowed_comm_ffi", out_type, has_side_effect=True)(
+            x, replica_groups=rg, group_size=gs, communication_id=cid
+        )
+
+    dummy = jnp.zeros((world,), dtype=jnp.uint8)
+    fn = jax.jit(shard_map(_body, mesh=mesh, in_specs=spec, out_specs=spec))
+    jax.block_until_ready(fn(dummy))
+
+
 def set_ep_config(config: EpConfig) -> None:
     """Cache the EP config for abstract-eval / sharding helpers. Call once."""
     global _ep_config

transformer_engine/jax/csrc/extensions.h

@@ -210,7 +210,7 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedMoEAuxLossBackwardHandler);
 void SetEpBootstrapParams(pybind11::bytes unique_id_bytes, int ep_size, int rank_within_group,
                           int num_experts, int max_tokens_per_rank, int max_recv_tokens_per_rank,
                           int hidden_dim, int max_num_sms, int max_token_dtype,
-                          bool drop_on_overflow);
+                          bool drop_on_overflow, bool borrowed_comm);
 void ReleaseEpResources();
 // Return the handle_mem byte size for a layer config.
 size_t EpHandleMemSize(int top_k, size_t dispatch_output_per_expert_alignment);
@@ -227,6 +227,11 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(EpCombineHandler);
 XLA_FFI_DECLARE_HANDLER_SYMBOL(EpDispatchBwdHandler);
 XLA_FFI_DECLARE_HANDLER_SYMBOL(EpCombineBwdHandler);
 
+// EP-specific execute stage of the borrowed-comm bootstrap op (see
+// tex.ep.use_nccl_comm_from_xla). The prepare stage is the generic
+// FfiRequestCliqueHandler in extensions/ffi_collectives.h.
+XLA_FFI_DECLARE_HANDLER_SYMBOL(EpBootstrapBorrowedCommHandler);
+
 // TopK
 XLA_FFI_DECLARE_HANDLER_SYMBOL(TopkHandler);
 pybind11::tuple GetTopkWorkspaceSizes(int batch_size, int seq_len, int k);