[JAX] Schedule EP dispatch/combine on XLA collective stream

tests/jax/test_multi_process_ep.py

@@ -36,6 +36,7 @@
     ep_combine_fwd,
     get_ep_config,
 )
+from transformer_engine.jax.version_utils import is_collective_stream_supported
 
 
 # ── Test config ─────────────────────────────────────────────────────────────
@@ -660,6 +661,55 @@ def run(idx, toks, w):
             expected = (("dp", "ep"),) if self.dp > 1 else ("ep",)
             self.assertEqual(tuple(compiled.output_shardings.spec), expected)
 
+    @unittest.skipUnless(
+        is_collective_stream_supported(),
+        "JAX/XLA lacks the gpu_stream:collective annotation (openxla/xla#39604)",
+    )
+    def test_z_dispatch_combine_on_collective_stream(self):
+        """Every EP FFI custom call must carry the collective-stream annotation
+        so XLA schedules them on the collective stream instead of overlapping
+        them with other collectives."""
+        T_dp, tokens, topk_idx, topk_w = self._make_random_inputs()
+        dp_spec = PartitionSpec(("dp", "ep"), None)
+        ep_spec_3d = PartitionSpec(("dp", "ep"), None, None)
+        ep_spec_2d = PartitionSpec(("dp", "ep"), None)
+
+        with self.mesh, global_shard_guard(self.mr):
+
+            @jax.jit
+            def run(idx, toks, w):
+                idx = jax.lax.with_sharding_constraint(idx, NamedSharding(self.mesh, dp_spec))
+                toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, dp_spec))
+                w = jax.lax.with_sharding_constraint(w, NamedSharding(self.mesh, dp_spec))
+                recv_t, recv_w, hm, tc = ep_dispatch(
+                    self.hk, idx, toks, w, self.recv_capacity_per_rank
+                )
+                recv_t = jax.lax.with_sharding_constraint(
+                    recv_t, NamedSharding(self.mesh, ep_spec_3d)
+                )
+                recv_w = jax.lax.with_sharding_constraint(
+                    recv_w, NamedSharding(self.mesh, ep_spec_2d)
+                )
+                weighted = self._preweight_expert_out(recv_t, recv_w)
+                out = ep_combine(self.hk, hm, tc, weighted, T_dp, out_sharding=(("dp", "ep"), None))
+                return jax.lax.with_sharding_constraint(out, NamedSharding(self.mesh, dp_spec))
+
+            hlo = run.lower(topk_idx, tokens, topk_w).compile().as_text()
+
+        # Every te_ep_* FFI custom call must carry the collective-stream
+        # annotation so XLA places it on the collective stream.
+        ep_lines = [l for l in hlo.splitlines() if 'custom_call_target="te_ep_' in l]
+        self.assertTrue(ep_lines, f"no te_ep_* custom calls in compiled HLO:\n{hlo}")
+        missing = [
+            l.strip()[:200]
+            for l in ep_lines
+            if '_xla_stream_annotation="collective"' not in l.replace(" ", "")
+        ]
+        self.assertFalse(
+            missing,
+            "te_ep_* custom calls missing collective-stream annotation:\n" + "\n".join(missing),
+        )
+
     def test_z_no_unexpected_reshard_in_hlo_bwd(self):
         """Compiled bwd HLO must not insert XLA collectives outside the EP FFI."""
         T_dp, tokens, topk_idx, topk_w = self._make_random_inputs()

transformer_engine/jax/cpp_extensions/ep.py

@@ -24,6 +24,18 @@
 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
+
+
+def _on_collective_stream(func):
+    """Pin ``func``'s ops to XLA's collective stream so the scheduler serializes
+    them with native collectives. No-op on JAX that lacks the annotation."""
+    if not is_collective_stream_supported():
+        return func
+    from jax.experimental.compute_on import compute_on
+
+    return compute_on("gpu_stream:collective")(func)  # pylint: disable=not-callable
+
 
 __all__ = [
     "EpConfig",
@@ -894,6 +906,7 @@ def shardy_sharding_rule(*args):
 # ── Public-ish helpers (used by jax/ep.py) ──────────────────────────────────
 
 
+@_on_collective_stream
 def ep_prepare(cfg: EpLayerConfig, topk_idx):
     """Exchange routing metadata for ``cfg``; return ``(token_counts, handle_mem)``."""
     return EpPreparePrimitive.outer_primitive.bind(
@@ -904,6 +917,7 @@ def ep_prepare(cfg: EpLayerConfig, topk_idx):
     )
 
 
+@_on_collective_stream
 def ep_dispatch_fwd(
     cfg: EpLayerConfig, handle_mem, topk_idx, tokens, topk_weights, recv_capacity_per_rank
 ):
@@ -920,6 +934,7 @@ def ep_dispatch_fwd(
     )
 
 
+@_on_collective_stream
 def ep_combine_fwd(
     cfg: EpLayerConfig, handle_mem, expert_out, num_local_tokens, out_partition_spec=None
 ):
@@ -935,6 +950,7 @@ def ep_combine_fwd(
     )
 
 
+@_on_collective_stream
 def ep_dispatch_bwd(
     cfg: EpLayerConfig,
     handle_mem,
@@ -956,6 +972,7 @@ def ep_dispatch_bwd(
     )
 
 
+@_on_collective_stream
 def ep_combine_bwd(cfg: EpLayerConfig, handle_mem, grad, recv_capacity_per_rank):
     """Backward of combine; returns grad_expert_out [num_procs, recv_capacity_per_rank, H]."""
     return EpCombineBwdPrimitive.outer_primitive.bind(

transformer_engine/jax/version_utils.py

@@ -64,6 +64,16 @@ def is_triton_autotuned_alias_safe() -> bool:
     return v >= PkgVersion(_TRITON_AUTOTUNED_ALIAS_STABLE_FLOOR)
 
 
+# XLA gained the ``gpu_stream:collective`` stream annotation in openxla/xla#39604,
+# which ships in the JAX 0.10.0 release. Older XLA fatally fails on it.
+_COLLECTIVE_STREAM_MIN_JAX_VERSION = "0.10.0"
+
+
+def is_collective_stream_supported() -> bool:
+    """Return True if the installed JAX supports the gpu_stream:collective annotation."""
+    return jax_version_meet_requirement(_COLLECTIVE_STREAM_MIN_JAX_VERSION)
+
+
 def is_triton_extension_supported() -> bool:
     """Return True if the current JAX version supports Triton kernel dispatch.
 
@@ -77,6 +87,7 @@ def is_triton_extension_supported() -> bool:
 __all__ = [
     "jax_version_meet_requirement",
     "is_triton_autotuned_alias_safe",
+    "is_collective_stream_supported",
     "is_triton_extension_supported",
     "TRITON_EXTENSION_MIN_JAX_VERSION",
     "TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION",