[PyTorch] [torch.compile] Prepare LayerNormLinear and LayerNormMLP for torch.compile

build_tools/VERSION.txt

@@ -1 +1 @@
-2.20.0.dev0
+2.21.0-dev0

build_tools/jax.py

@@ -5,6 +5,7 @@
 """JAX related extensions."""
 
 import os
+import warnings
 from pathlib import Path
 from packaging import version
 
@@ -95,16 +96,24 @@ def setup_jax_extension(
     if (discovered_nccl_include_path := nccl_include_path()) is not None:
         include_dirs.append(discovered_nccl_include_path)
     include_dirs.append(cudnn_frontend_include_path())
+    xla_include_path = xla_path()
     include_dirs.extend(
         [
             common_header_files,
             common_header_files / "common",
             common_header_files / "common" / "include",
             csrc_header_files,
-            xla_path(),
+            xla_include_path,
         ]
     )
 
+    # Match the borrowed-comm path's compile-time header check.
+    if not (Path(xla_include_path) / "xla/ffi/api/collectives_c_api.h").is_file():
+        warnings.warn(
+            f"XLA headers in {xla_include_path} do not include "
+            "xla/ffi/api/collectives_c_api.h; 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/multi_process_launch_ep.sh

@@ -7,6 +7,16 @@
 SCRIPT_NAMES="${SCRIPT_NAMES:-test_multi_process_ep.py}"
 TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-180}"
 
+# Each communicator mode needs a fresh process group.
+if [ -z "${NVTE_TEST_EP_CLASSES:-}" ]; then
+  RET=0
+  NVTE_TEST_EP_CLASSES="TestEP,TestEPOverflowDrop,TestEpDomainGrouping" \
+    bash "${BASH_SOURCE[0]}" || RET=1
+  NVTE_TEST_EP_CLASSES="TestEPBorrowedComm" \
+    bash "${BASH_SOURCE[0]}" || RET=1
+  exit "$RET"
+fi
+
 
 XLA_BASE_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
                 --xla_gpu_graph_min_graph_size=1"

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 ─────────────────────────────────────────────────────────
 
 
@@ -846,6 +904,8 @@ 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})")
+        cls._prev_comm_env = os.environ.get("NVTE_JAX_EP_NCCL_COMM_FROM_XLA")
+        os.environ["NVTE_JAX_EP_NCCL_COMM_FROM_XLA"] = "0"
         cls.num_procs = jax.process_count()
         cls.rank = jax.process_index()
         cls.dp, cls.ep = _factor_dp_ep(cls.num_procs)
@@ -873,6 +933,10 @@ def setUpClass(cls):
     def tearDownClass(cls):
         # Leave a clean slate so another class can bootstrap after us.
         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
 
     def _make_concentrated_inputs(self):
         """All top-1 routes to expert 0; top-2 spread over the rest, so the rank
@@ -961,6 +1025,65 @@ 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)
+
+
+def _ep_test_cases():
+    """Select test classes for one communicator mode."""
+    all_test_cases = {
+        c.__name__: c
+        for c in (TestEP, TestEPBorrowedComm, TestEPOverflowDrop, TestEpDomainGrouping)
+    }
+    names = os.environ.get("NVTE_TEST_EP_CLASSES")
+    test_cases = (
+        tuple(all_test_cases[name.strip()] for name in names.split(","))
+        if names
+        else (TestEP, TestEPOverflowDrop, TestEpDomainGrouping)
+    )
+    if TestEPBorrowedComm in test_cases and any(
+        c in test_cases for c in (TestEP, TestEPOverflowDrop)
+    ):
+        raise ValueError("Run borrowed-comm and self-hosted EP tests in separate processes.")
+    return test_cases
+
+
 # ── Entry point ──────────────────────────────────────────────────────────────
 
 
@@ -972,6 +1095,20 @@ def test_ep_tp_splits_domains(self):
     coord_addr = sys.argv[1]
     proc_id = int(sys.argv[2])
     num_procs = int(sys.argv[3])
+    test_cases = _ep_test_cases()
+
+    target = os.environ.get("TARGET_TEST")
+    if target:
+        name = target.split(".")[-1]
+        if not any(
+            hasattr(c, name)
+            for c in (TestEP, TestEPBorrowedComm, TestEPOverflowDrop, TestEpDomainGrouping)
+        ):
+            raise ValueError(f"Unknown EP test: {target}")
+        test_cases = tuple(c for c in test_cases if hasattr(c, name))
+        if not test_cases:
+            unittest.TextTestRunner(verbosity=2).run(unittest.TestSuite())
+            sys.exit(0)
 
     jax.distributed.initialize(
         coordinator_address=coord_addr,
@@ -981,12 +1118,8 @@ def test_ep_tp_splits_domains(self):
     )
 
     loader = unittest.TestLoader()
-    test_cases = (TestEP, TestEPOverflowDrop, TestEpDomainGrouping)
-    target = os.environ.get("TARGET_TEST")
     if target:
-        name = target.split(".")[-1]
-        cls = next((c for c in test_cases if hasattr(c, name)), TestEP)
-        suite = loader.loadTestsFromName(name, cls)
+        suite = unittest.TestSuite(loader.loadTestsFromName(name, c) for c in test_cases)
     else:
         suite = unittest.TestSuite(loader.loadTestsFromTestCase(c) for c in test_cases)
     runner = unittest.TextTestRunner(verbosity=2)

tests/pytorch/distributed/test_comm_gemm_overlap.py

@@ -33,6 +33,15 @@
 # to avoid numerical tolerance issues of doing comm gemm overlap, limit the number of GPUs used
 MAX_GPUS_TO_USE = 4
 
+COMM_GEMM_QUANTIZATION_PARAMS = [
+    pytest.param(False, "none", id="ub-bf16"),
+    pytest.param(False, "fp8", id="ub-fp8"),
+    pytest.param(False, "mxfp8", id="ub-mxfp8"),
+    pytest.param(True, "none", id="cublasmp-bf16"),
+    pytest.param(True, "fp8", id="cublasmp-fp8"),
+    pytest.param(True, "mxfp8", id="cublasmp-mxfp8"),
+]
+
 TEST_ROOT = Path(__file__).parent.resolve()
 NUM_PROCS: int = min(torch.cuda.device_count(), MAX_GPUS_TO_USE)
 LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"]
@@ -102,10 +111,6 @@ def _run_gemm_with_overlap(
         if use_cublasmp:
             if not tex.nvte_built_with_cublasmp():
                 pytest.skip("Transformer Engine not built with cuBLASMp (NVTE_WITH_CUBLASMP=0).")
-            if quantization == "mxfp8":
-                pytest.skip(
-                    "cuBLASMp comm+GEMM overlap does not yet support MXFP8 (block scaling)."
-                )
             if comm_type == "RS" and not p2p and not tex.device_supports_multicast():
                 pytest.skip(
                     "cuBLASMp non-P2P reduce-scatter requires NVSwitch (multicast support)."
@@ -159,8 +164,6 @@ def _run_layer_with_overlap(
     if use_cublasmp:
         if not tex.nvte_built_with_cublasmp():
             pytest.skip("Transformer Engine not built with cuBLASMp (NVTE_WITH_CUBLASMP=0).")
-        if fp8 and quantization == "mxfp8":
-            pytest.skip("cuBLASMp comm+GEMM overlap does not yet support MXFP8 (block scaling).")
         test_cmd.append("--use-cublasmp")
 
     test_env = os.environ.copy()
@@ -182,8 +185,7 @@ def _run_layer_with_overlap(
     _assert_subprocess_succeeded(result)
 
 
-@pytest.mark.parametrize("use_cublasmp", (False, True))
-@pytest.mark.parametrize("quantization", ("none", "fp8", "mxfp8"))
+@pytest.mark.parametrize("use_cublasmp,quantization", COMM_GEMM_QUANTIZATION_PARAMS)
 @pytest.mark.parametrize("aggregate", (False, True))
 def test_split_all_gather_overlaps(quantization, aggregate, use_cublasmp):
     """
@@ -193,8 +195,7 @@ def test_split_all_gather_overlaps(quantization, aggregate, use_cublasmp):
     _run_gemm_with_overlap("AG", False, True, False, aggregate, quantization, use_cublasmp)
 
 
-@pytest.mark.parametrize("use_cublasmp", (False, True))
-@pytest.mark.parametrize("quantization", ("none", "fp8", "mxfp8"))
+@pytest.mark.parametrize("use_cublasmp,quantization", COMM_GEMM_QUANTIZATION_PARAMS)
 @pytest.mark.parametrize("p2p", (False, True))
 def test_split_reduce_scatter_overlaps(quantization, p2p, use_cublasmp):
     """