diff --git a/build_tools/jax.py b/build_tools/jax.py index d61c26c128..445305a165 100644 --- a/build_tools/jax.py +++ b/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(): diff --git a/qa/L2_jax_distributed_unittest/test.sh b/qa/L2_jax_distributed_unittest/test.sh index 330b254e7d..c64a1e561d 100644 --- a/qa/L2_jax_distributed_unittest/test.sh +++ b/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 diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 47af0b0c39..21842e299f 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/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] diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index ca70ea145c..ca22484a26 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/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 diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 580219baf2..ea8793423a 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/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); diff --git a/transformer_engine/jax/csrc/extensions/ep.cpp b/transformer_engine/jax/csrc/extensions/ep.cpp index aa5ed27faa..ec9bab4769 100644 --- a/transformer_engine/jax/csrc/extensions/ep.cpp +++ b/transformer_engine/jax/csrc/extensions/ep.cpp @@ -18,6 +18,7 @@ #include "../extensions.h" #include "common.h" +#include "ffi_collectives.h" #include "transformer_engine/gemm.h" namespace transformer_engine { @@ -36,26 +37,32 @@ struct EpBootstrapParams { int max_num_sms = 0; NVTEDType max_token_dtype = kNVTEBFloat16; bool drop_on_overflow = false; + // When set, EP borrows XLA's comm (see below): no ncclCommInitRank at + // bootstrap; nvte_ep_initialize is deferred to the first executable that + // fetches the borrowed communicator. + bool borrowed_comm = false; }; +static NVTEEpGroupConfig MakeEpGroupConfig(const EpBootstrapParams& p) { + return NVTEEpGroupConfig{.struct_size = sizeof(NVTEEpGroupConfig), + .ep_size = p.ep_size, + .num_experts = p.num_experts, + .max_tokens_per_rank = p.max_tokens_per_rank, + .max_recv_tokens_per_rank = p.max_recv_tokens_per_rank, + .hidden_dim = p.hidden_dim, + .num_comm_sms = p.max_num_sms, + .max_token_dtype = p.max_token_dtype, + .zero_copy = 0, + .drop_on_overflow = p.drop_on_overflow}; +} + class EpResources { public: explicit EpResources(const EpBootstrapParams& p) { ncclUniqueId uid; std::memcpy(&uid, p.uid_bytes.data(), sizeof(uid)); NVTE_CHECK_NCCL(ncclCommInitRank(&comm_, p.ep_size, uid, p.rank_within_group)); - // zero_copy=0: JAX EP path always stages payloads; the zero-copy fast path - // requires NVTECommWindow-backed tensors, which JAX bindings don't expose. - NVTEEpGroupConfig cfg{.struct_size = sizeof(NVTEEpGroupConfig), - .ep_size = p.ep_size, - .num_experts = p.num_experts, - .max_tokens_per_rank = p.max_tokens_per_rank, - .max_recv_tokens_per_rank = p.max_recv_tokens_per_rank, - .hidden_dim = p.hidden_dim, - .num_comm_sms = p.max_num_sms, - .max_token_dtype = p.max_token_dtype, - .zero_copy = 0, - .drop_on_overflow = p.drop_on_overflow}; + NVTEEpGroupConfig cfg = MakeEpGroupConfig(p); try { nvte_ep_initialize(static_cast(comm_), &cfg); } catch (...) { @@ -97,6 +104,21 @@ bool g_ep_params_set = false; std::weak_ptr g_ep_resources_weak; // Python-held anchor so trace-time handle_mem allocs find EPBackend ready. std::shared_ptr g_ep_resources_anchor; +// Borrowed-comm path: EPBackend is initialized once from a borrowed communicator. +bool g_ep_xla_initialized = false; + +#ifdef NVTE_FFI_COLLECTIVES_AVAILABLE +// Idempotently initialize EPBackend on a borrowed communicator. Safe to call +// from every executable that fetches the comm; only the first call initializes. +void EnsureEpBackendFromBorrowedComm(ncclComm_t comm) { + std::lock_guard lock(g_ep_mu); + if (g_ep_xla_initialized) return; + NVTE_CHECK(g_ep_params_set, "EP bootstrap params not set before borrowing XLA comm."); + NVTEEpGroupConfig cfg = MakeEpGroupConfig(g_ep_params); + nvte_ep_initialize(static_cast(comm), &cfg); + g_ep_xla_initialized = true; +} +#endif // collectives header available std::shared_ptr AcquireEpResources() { std::lock_guard lock(g_ep_mu); @@ -128,14 +150,14 @@ struct EpConfig { void SetEpBootstrapParams(pybind11::bytes unique_id_bytes_obj, 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) { std::string uid_str = unique_id_bytes_obj; NVTE_CHECK(static_cast(uid_str.size()) >= 128, "unique_id_bytes must be at least 128 bytes (ncclUniqueId size)."); std::shared_ptr anchor; { std::lock_guard lock(g_ep_mu); - NVTE_CHECK(!g_ep_resources_anchor, + NVTE_CHECK(!g_ep_resources_anchor && !g_ep_xla_initialized, "EP bootstrap already initialized; call release_ep_resources() before re-init."); std::memcpy(g_ep_params.uid_bytes.data(), uid_str.data(), 128); g_ep_params.ep_size = ep_size; @@ -147,8 +169,11 @@ void SetEpBootstrapParams(pybind11::bytes unique_id_bytes_obj, int ep_size, int g_ep_params.max_num_sms = max_num_sms; g_ep_params.max_token_dtype = static_cast(max_token_dtype); g_ep_params.drop_on_overflow = drop_on_overflow; + g_ep_params.borrowed_comm = borrowed_comm; g_ep_params_set = true; } + // Borrowed-comm path defers NCCL init to the first executable; nothing eager. + if (borrowed_comm) return; // Acquire outside the lock: EpResources ctor runs ncclCommInitRank which is // a collective and may block on peer ranks. anchor = AcquireEpResources(); @@ -157,11 +182,17 @@ void SetEpBootstrapParams(pybind11::bytes unique_id_bytes_obj, int ep_size, int } // Drops the anchor; comm tears down once the last executable also releases. +// For the borrowed-comm path, tears down EPBackend while the borrowed comm is +// still alive (call from ep_finalize, not atexit). void ReleaseEpResources() { std::shared_ptr to_drop; { std::lock_guard lock(g_ep_mu); to_drop = std::move(g_ep_resources_anchor); + if (g_ep_xla_initialized) { + nvte_ep_shutdown(); + g_ep_xla_initialized = false; + } } // to_drop dtor runs outside the lock. } @@ -186,6 +217,12 @@ pybind11::capsule GetEpInstanceStateTypeInfoCapsule() { static ::xla::ffi::ErrorOr> EpInstantiateImpl() { auto state = std::make_unique(); + { + // Borrowed-comm path: XLA owns the comm, so there is nothing self-hosted to + // acquire or pin per-executable; EPBackend is initialized by the bootstrap op. + std::lock_guard lock(g_ep_mu); + if (g_ep_params.borrowed_comm) return state; + } try { state->resources = AcquireEpResources(); } catch (const std::exception& e) { @@ -479,6 +516,53 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpCombineBwdHandler, EpCombineBwdFFI, .Attrs(), FFI_CudaGraph_Traits); +// -- Borrowed-comm path ------------------------------------------------------- +// +// Instead of a self-hosted ncclCommInitRank, EPBackend borrows the communicator +// XLA already owns for the EP replica groups. A one-shot bootstrap op fetches it +// and initializes the backend once; the per-step ops then share the same FFI +// targets as the self-hosted path (see EpInstantiateImpl). Auto-selected from +// Python (see tex.ep.use_nccl_comm_from_xla). The prepare stage is the generic +// FfiRequestCliqueHandler (see ffi_collectives.cpp); only the EP-specific +// execute stage lives here. +#ifdef NVTE_FFI_COLLECTIVES_AVAILABLE + +// Execute stage: fetch the borrowed comm and initialize EPBackend once. The +// token buffer is copied straight through only to give the op an input->output +// data dependency, so it stays ordered on the borrowed comm's stream. +Error_Type EpBootstrapBorrowedCommFFI(cudaStream_t stream, EpInstanceState* ep_state, + FfiCollectivesCtx coll, Buffer_Type token, Result_Type out, + Span_Type replica_groups, int64_t group_size, + int64_t communication_id) { + (void)ep_state; + auto groups = ffi_collectives::ReplicaGroupsFromFlat(replica_groups.begin(), + replica_groups.size(), group_size); + auto comm_or = ffi_collectives::GetComm(coll, groups, communication_id); + if (comm_or.has_error()) return comm_or.error(); + ncclComm_t comm = comm_or.value(); + NVTE_CHECK(comm != nullptr, "XLA returned a null EP communicator."); + EnsureEpBackendFromBorrowedComm(comm); + const size_t bytes = token.size_bytes(); + if (bytes > 0) { + NVTE_CHECK_CUDA(cudaMemcpyAsync(out->untyped_data(), token.untyped_data(), bytes, + cudaMemcpyDeviceToDevice, stream)); + } + return ffi_with_cuda_error_check(); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(EpBootstrapBorrowedCommHandler, EpBootstrapBorrowedCommFFI, + FFI::Bind() + .Ctx() // stream + .Ctx<::xla::ffi::State>() // EP state + .Ctx<::xla::ffi::Extension>() + .Arg() // token (identity in) + .Ret() // token (identity out) + .Attr>("replica_groups") + .Attr("group_size") + .Attr("communication_id")); + +#endif // collectives header available + } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/ffi.h b/transformer_engine/jax/csrc/extensions/ffi.h index f9d327102b..6a4b55bfc7 100644 --- a/transformer_engine/jax/csrc/extensions/ffi.h +++ b/transformer_engine/jax/csrc/extensions/ffi.h @@ -19,6 +19,8 @@ using Result_Type = xla::ffi::Result; using Variadic_Buffer_Type = xla::ffi::RemainingArgs; using Variadic_Result_Type = xla::ffi::RemainingRets; using Error_Type = xla::ffi::Error; +template +using Span_Type = xla::ffi::Span; using FFI = xla::ffi::Ffi; using FFI_Stream_Type = xla::ffi::PlatformStream; using Dictionary = xla::ffi::Dictionary; diff --git a/transformer_engine/jax/csrc/extensions/ffi_collectives.cpp b/transformer_engine/jax/csrc/extensions/ffi_collectives.cpp new file mode 100644 index 0000000000..38087aee0d --- /dev/null +++ b/transformer_engine/jax/csrc/extensions/ffi_collectives.cpp @@ -0,0 +1,103 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "ffi_collectives.h" + +#ifdef NVTE_FFI_COLLECTIVES_AVAILABLE + +#include + +#include "ffi.h" + +namespace transformer_engine { +namespace jax { + +namespace ffi_collectives { + +namespace { + +::xla::ffi::Error TakeError(const XLA_FFI_Api* api, XLA_FFI_Error* err) { + std::string msg = ::xla::ffi::internal::GetErrorMessage(api, err); + ::xla::ffi::internal::DestroyError(api, err); + return ::xla::ffi::Error::Internal(msg); +} + +} // namespace + +std::vector ToRawGroups(const std::vector>& groups) { + std::vector raw; + raw.reserve(groups.size()); + for (const auto& g : groups) { + raw.push_back(XLA_FFI_ReplicaGroup{g.data(), g.size()}); + } + return raw; +} + +::xla::ffi::Error RequestClique(const FfiCollectivesCtx& ctx, + const std::vector>& groups, + int64_t communication_id) { + std::vector raw = ToRawGroups(groups); + XLA_FFI_Communicator_Request_Args args; + args.struct_size = XLA_FFI_Communicator_Request_Args_STRUCT_SIZE; + args.extension_start = nullptr; + args.group_mode = XLA_FFI_GROUP_FLATTENED_ID; + args.groups = raw.data(); + args.num_groups = raw.size(); + args.communication_id = communication_id; + if (XLA_FFI_Error* err = ctx.ext->request_communicator(ctx.ext, &args)) { + return TakeError(ctx.api, err); + } + return ::xla::ffi::Error::Success(); +} + +::xla::ffi::ErrorOr GetComm(const FfiCollectivesCtx& ctx, + const std::vector>& groups, + int64_t communication_id) { + std::vector raw = ToRawGroups(groups); + XLA_FFI_Communicator_Get_Args args; + args.struct_size = XLA_FFI_Communicator_Get_Args_STRUCT_SIZE; + args.extension_start = nullptr; + args.group_mode = XLA_FFI_GROUP_FLATTENED_ID; + args.groups = raw.data(); + args.num_groups = raw.size(); + args.communication_id = communication_id; + args.communicator = nullptr; + if (XLA_FFI_Error* err = ctx.ext->get_communicator(ctx.ext, &args)) { + return TakeError(ctx.api, err); + } + return reinterpret_cast(args.communicator); +} + +std::vector> ReplicaGroupsFromFlat(const int64_t* flat, size_t count, + int64_t group_size) { + std::vector> groups; + if (group_size <= 0) return groups; + for (size_t off = 0; off + group_size <= count; off += group_size) { + groups.emplace_back(flat + off, flat + off + group_size); + } + return groups; +} + +} // namespace ffi_collectives + +Error_Type FfiRequestCliqueFFI(FfiCollectivesCtx coll, Span_Type replica_groups, + int64_t group_size, int64_t communication_id) { + auto groups = ffi_collectives::ReplicaGroupsFromFlat(replica_groups.begin(), + replica_groups.size(), group_size); + return ffi_collectives::RequestClique(coll, groups, communication_id); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(FfiRequestCliqueHandler, FfiRequestCliqueFFI, + FFI::BindPrepare() + .Ctx<::xla::ffi::Extension>() + .Attr>("replica_groups") + .Attr("group_size") + .Attr("communication_id")); + +} // namespace jax +} // namespace transformer_engine + +#endif // collectives header available diff --git a/transformer_engine/jax/csrc/extensions/ffi_collectives.h b/transformer_engine/jax/csrc/extensions/ffi_collectives.h new file mode 100644 index 0000000000..fb25d07a03 --- /dev/null +++ b/transformer_engine/jax/csrc/extensions/ffi_collectives.h @@ -0,0 +1,86 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file ffi_collectives.h + * \brief Borrow the XLA-owned NCCL communicator inside any FFI handler. + * + * Not EP-specific -- any FFI handler that wants XLA's comm can request the + * clique (prepare stage) and fetch the borrowed ncclComm_t (execute stage). + * Absent on older XLA, in which case the borrow path must not be selected. + */ + +#ifndef TRANSFORMER_ENGINE_JAX_CSRC_EXTENSIONS_FFI_COLLECTIVES_H_ +#define TRANSFORMER_ENGINE_JAX_CSRC_EXTENSIONS_FFI_COLLECTIVES_H_ + +// NVTE_FFI_COLLECTIVES_AVAILABLE is the single source of truth for "is the +// XLA collectives FFI extension available"; callers add their own build gates +// (e.g. NVTE_WITH_NCCL_EP) on top of this. +#if __has_include("xla/ffi/api/collectives_c_api.h") +#define NVTE_FFI_COLLECTIVES_AVAILABLE 1 + +#include + +#include +#include + +#include "xla/ffi/api/collectives_c_api.h" +#include "xla/ffi/api/ffi.h" + +namespace transformer_engine { +namespace jax { + +// Decoded context: the FFI api table plus the found collectives extension. +struct FfiCollectivesCtx { + const XLA_FFI_Api* api = nullptr; + const XLA_FFI_Collectives_Extension* ext = nullptr; +}; + +// Trait type for ::xla::ffi::Extension. The public FFI +// CtxDecoding looks the extension up by kExtensionType and hands us a typed +// context, so we do not depend on XLA-internal headers that jaxlib omits. +struct FfiCollectives { + using Type = FfiCollectivesCtx; + using CExtension = XLA_FFI_Collectives_Extension; + static constexpr const char* kName = "CollectivesExtension"; + static constexpr int32_t kExtensionType = XLA_FFI_Extension_Collectives; + static constexpr int32_t kMajorVersion = XLA_FFI_Extension_Collectives_MajorVersion; + static constexpr int32_t kMinorVersion = XLA_FFI_Extension_Collectives_MinorVersion; + // Accept any minor within the same major so a newer runtime still binds. + static bool Support(int32_t major, int32_t /*minor*/) { return major == kMajorVersion; } + static Type Create(const XLA_FFI_Api* api, const CExtension* ext) { return Type{api, ext}; } +}; + +namespace ffi_collectives { + +std::vector ToRawGroups(const std::vector>& groups); + +// Prepare stage: ask XLA to acquire the clique for `groups` (flattened-id mode). +::xla::ffi::Error RequestClique(const FfiCollectivesCtx& ctx, + const std::vector>& groups, + int64_t communication_id); + +// Execute stage: fetch the borrowed communicator (ncclComm_t on XLA:GPU). +::xla::ffi::ErrorOr GetComm(const FfiCollectivesCtx& ctx, + const std::vector>& groups, + int64_t communication_id); + +// Rebuild ragged replica groups from a flat buffer of equal-size groups. +std::vector> ReplicaGroupsFromFlat(const int64_t* flat, size_t count, + int64_t group_size); + +} // namespace ffi_collectives + +// Generic prepare-stage handler: any FFI op that borrows XLA's comm binds this +// to request the clique before execute. Reads int64 attrs "replica_groups" +// (flat, equal-size groups), "group_size", and "communication_id". +XLA_FFI_DECLARE_HANDLER_SYMBOL(FfiRequestCliqueHandler); + +} // namespace jax +} // namespace transformer_engine + +#endif // collectives header available + +#endif // TRANSFORMER_ENGINE_JAX_CSRC_EXTENSIONS_FFI_COLLECTIVES_H_ diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 3927e2686e..014024d9bb 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -7,6 +7,7 @@ #include "../extensions.h" #include "cgemm_helper.h" #include "common/util/cuda_runtime.h" +#include "ffi_collectives.h" // FfiRequestCliqueHandler (borrowed-comm prepare) #include "transformer_engine/gemm.h" namespace transformer_engine { @@ -124,6 +125,17 @@ pybind11::dict Registrations() { dict["te_ep_combine_bwd_ffi"] = pybind11::dict(pybind11::arg("instantiate") = EncapsulateFFI(EpInstantiateHandler), pybind11::arg("execute") = EncapsulateFFI(EpCombineBwdHandler)); + + // Borrowed-comm bootstrap: a one-shot op requests the collective clique + // (generic prepare) and initializes EPBackend from the borrowed comm (EP + // execute). Registered only when the XLA collectives headers were available + // at build time; its absence is how Python detects an unbuilt path. +#ifdef NVTE_FFI_COLLECTIVES_AVAILABLE + dict["te_ep_bootstrap_borrowed_comm_ffi"] = + pybind11::dict(pybind11::arg("instantiate") = EncapsulateFFI(EpInstantiateHandler), + pybind11::arg("prepare") = EncapsulateFFI(FfiRequestCliqueHandler), + pybind11::arg("execute") = EncapsulateFFI(EpBootstrapBorrowedCommHandler)); +#endif // collectives header available #endif // NVTE_WITH_NCCL_EP // TopK @@ -160,7 +172,7 @@ PYBIND11_MODULE(transformer_engine_jax, m) { pybind11::arg("ep_size"), pybind11::arg("rank_within_group"), pybind11::arg("num_experts"), pybind11::arg("max_tokens_per_rank"), pybind11::arg("max_recv_tokens_per_rank"), pybind11::arg("hidden_dim"), pybind11::arg("max_num_sms"), pybind11::arg("max_token_dtype"), - pybind11::arg("drop_on_overflow")); + pybind11::arg("drop_on_overflow"), pybind11::arg("borrowed_comm") = false); m.def("release_ep_resources", &ReleaseEpResources); m.def("ep_handle_mem_size", &EpHandleMemSize, pybind11::arg("top_k"), pybind11::arg("dispatch_output_per_expert_alignment") = 0); diff --git a/transformer_engine/jax/ep.py b/transformer_engine/jax/ep.py index 2222a41e48..ea74b86ce3 100644 --- a/transformer_engine/jax/ep.py +++ b/transformer_engine/jax/ep.py @@ -100,6 +100,21 @@ def device_to_rank(d): return int(grid[row, 0]), col, int(grid.shape[0]) +def _ep_flattened_replica_groups(mesh, ep_resource): + """FLATTENED_ID replica groups for the EP axis, as a flat int64 array. + + Each group fixes all non-ep mesh coordinates and varies ep. Returns + ``(flat_groups, ep_size)``. + """ + shape = tuple(mesh.shape[a] for a in mesh.axis_names) + ep_pos = mesh.axis_names.index(ep_resource) + ep_size = shape[ep_pos] + world = int(np.prod(shape)) + grid = np.arange(world, dtype=np.int64).reshape(shape) + groups = np.moveaxis(grid, ep_pos, -1).reshape(-1, ep_size) + return groups.reshape(-1), ep_size + + def ep_bootstrap( world_size, rank, @@ -175,6 +190,42 @@ def ep_bootstrap( if num_experts % ep_size != 0: raise ValueError(f"num_experts ({num_experts}) must be divisible by ep_size ({ep_size}).") + common_cfg = { + "world_size": world_size, + "rank": rank, + "ep_size": ep_size, + "num_ep_groups": num_ep_groups, + "num_experts": num_experts, + "num_local_experts": num_experts // ep_size, + "max_tokens_per_rank": max_tokens_per_rank, + "recv_capacity_per_rank": recv_capacity_per_rank, + "hidden_dim": hidden_dim, + } + + # Borrowed-comm path (auto-selected by tex.ep.use_nccl_comm_from_xla): XLA + # owns the EP communicator, so a one-shot bootstrap op fetches it and + # initializes EPBackend instead of a host-side UID exchange. + if tex.ep.use_nccl_comm_from_xla(): + replica_groups, ep_group_size = _ep_flattened_replica_groups(mesh, ep_resource) + communication_id = tex.ep.EP_COMMUNICATION_ID + transformer_engine_jax.set_ep_bootstrap_params( + bytes(128), + ep_size, + 0, + num_experts, + max_tokens_per_rank, + recv_capacity_per_rank, + hidden_dim, + max_num_sms=int(max_num_sms), + max_token_dtype=int(jax_dtype_to_te_dtype(max_token_dtype)), + drop_on_overflow=bool(drop_on_overflow), + borrowed_comm=True, + ) + tex.ep.set_ep_config(tex.ep.EpConfig(**common_cfg)) + # Initialize EPBackend now so trace-time handle_mem_size finds it ready. + tex.ep.run_borrowed_comm_bootstrap(mesh, replica_groups, ep_group_size, communication_id) + return + UID_SIZE = 128 root_rank, rank_within_group, _num_domains = _ep_domain_for_rank(mesh, ep_resource, rank) is_color_root = rank_within_group == 0 @@ -211,19 +262,7 @@ def ep_bootstrap( atexit.register(transformer_engine_jax.release_ep_resources) _atexit_registered = True - tex.ep.set_ep_config( - tex.ep.EpConfig( - world_size=world_size, - rank=rank, - ep_size=ep_size, - num_ep_groups=num_ep_groups, - num_experts=num_experts, - num_local_experts=num_experts // ep_size, - max_tokens_per_rank=max_tokens_per_rank, - recv_capacity_per_rank=recv_capacity_per_rank, - hidden_dim=hidden_dim, - ) - ) + tex.ep.set_ep_config(tex.ep.EpConfig(**common_cfg)) def ep_finalize(): diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py index 500c859b4c..18ee255f8f 100644 --- a/transformer_engine/jax/version_utils.py +++ b/transformer_engine/jax/version_utils.py @@ -84,6 +84,26 @@ def is_collective_stream_supported() -> bool: return True +# Minimum JAX version whose XLA ships the FFI collectives extension (lets an FFI +# handler fetch XLA's own communicator). Conservative floor: the first version +# this was verified on. +_XLA_FFI_COLLECTIVES_NIGHTLY_FLOOR = "0.11.2.dev20260828" +_XLA_FFI_COLLECTIVES_STABLE_FLOOR = "0.11.2" + + +@lru_cache(maxsize=None) +def is_xla_ffi_collectives_supported() -> bool: + """Return True if the installed JAX exposes the XLA FFI collectives extension. + + Not EP-specific; gates auto-selection of the EP borrowed-comm path (see + cpp_extensions.ep.use_nccl_comm_from_xla). + """ + v = PkgVersion(get_pkg_version("jax")) + if v.dev is not None: + return v >= PkgVersion(_XLA_FFI_COLLECTIVES_NIGHTLY_FLOOR) + return v >= PkgVersion(_XLA_FFI_COLLECTIVES_STABLE_FLOOR) + + def is_triton_extension_supported() -> bool: """Return True if the current JAX version supports Triton kernel dispatch. @@ -98,6 +118,7 @@ def is_triton_extension_supported() -> bool: "jax_version_meet_requirement", "is_triton_autotuned_alias_safe", "is_collective_stream_supported", + "is_xla_ffi_collectives_supported", "is_triton_extension_supported", "TRITON_EXTENSION_MIN_JAX_VERSION", "TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION",