From d4428a74869df74d9d44b1318627a34ce93f2d4b Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 21 Jul 2026 15:19:52 -0700 Subject: [PATCH 1/4] [JAX] Schedule EP dispatch/combine on XLA collective stream Signed-off-by: Phuong Nguyen --- tests/jax/test_multi_process_ep.py | 46 +++++++++++++++++++++ transformer_engine/jax/cpp_extensions/ep.py | 10 +++++ 2 files changed, 56 insertions(+) diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 0b8bb25f3f..67fdd4eb98 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -19,6 +19,7 @@ """ import os +import re import sys import unittest @@ -660,6 +661,51 @@ def run(idx, toks, w): expected = (("dp", "ep"),) if self.dp > 1 else ("ep",) self.assertEqual(tuple(compiled.output_shardings.spec), expected) + 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() diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 77e60afbcd..51f7c4da69 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -19,12 +19,17 @@ import jax import jax.numpy as jnp from jax import dtypes, ffi +from jax.experimental.compute_on import compute_on 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 +# Pin EP collectives to XLA's high-priority collective stream so the scheduler +# serializes them with native collectives instead of overlapping them. +_EP_COLLECTIVE_STREAM = "gpu_stream:collective" + __all__ = [ "EpConfig", "EpLayerConfig", @@ -894,6 +899,7 @@ def shardy_sharding_rule(*args): # ── Public-ish helpers (used by jax/ep.py) ────────────────────────────────── +@compute_on(_EP_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 +910,7 @@ def ep_prepare(cfg: EpLayerConfig, topk_idx): ) +@compute_on(_EP_COLLECTIVE_STREAM) def ep_dispatch_fwd( cfg: EpLayerConfig, handle_mem, topk_idx, tokens, topk_weights, recv_capacity_per_rank ): @@ -920,6 +927,7 @@ def ep_dispatch_fwd( ) +@compute_on(_EP_COLLECTIVE_STREAM) def ep_combine_fwd( cfg: EpLayerConfig, handle_mem, expert_out, num_local_tokens, out_partition_spec=None ): @@ -935,6 +943,7 @@ def ep_combine_fwd( ) +@compute_on(_EP_COLLECTIVE_STREAM) def ep_dispatch_bwd( cfg: EpLayerConfig, handle_mem, @@ -956,6 +965,7 @@ def ep_dispatch_bwd( ) +@compute_on(_EP_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( From 5d3271ec49e8e562b147ce80eb33e84a11a673f3 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 21 Jul 2026 15:28:19 -0700 Subject: [PATCH 2/4] Update tests/jax/test_multi_process_ep.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Phuong Nguyen --- tests/jax/test_multi_process_ep.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 67fdd4eb98..d24fcd8235 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -19,7 +19,6 @@ """ import os -import re import sys import unittest From 2b62e3c0c201ee032236bfaff9d368cf868d820c Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 21 Jul 2026 16:12:00 -0700 Subject: [PATCH 3/4] [JAX] Gate EP collective-stream annotation on JAX/XLA version Signed-off-by: Phuong Nguyen --- tests/jax/test_multi_process_ep.py | 5 +++++ transformer_engine/jax/cpp_extensions/ep.py | 25 +++++++++++++-------- transformer_engine/jax/version_utils.py | 11 +++++++++ 3 files changed, 32 insertions(+), 9 deletions(-) diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index d24fcd8235..d3effcdc07 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/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,10 @@ 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 diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 51f7c4da69..0723db6ebe 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -19,16 +19,23 @@ import jax import jax.numpy as jnp from jax import dtypes, ffi -from jax.experimental.compute_on import compute_on 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 + + +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) -# Pin EP collectives to XLA's high-priority collective stream so the scheduler -# serializes them with native collectives instead of overlapping them. -_EP_COLLECTIVE_STREAM = "gpu_stream:collective" __all__ = [ "EpConfig", @@ -899,7 +906,7 @@ def shardy_sharding_rule(*args): # ── Public-ish helpers (used by jax/ep.py) ────────────────────────────────── -@compute_on(_EP_COLLECTIVE_STREAM) +@_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( @@ -910,7 +917,7 @@ def ep_prepare(cfg: EpLayerConfig, topk_idx): ) -@compute_on(_EP_COLLECTIVE_STREAM) +@_on_collective_stream def ep_dispatch_fwd( cfg: EpLayerConfig, handle_mem, topk_idx, tokens, topk_weights, recv_capacity_per_rank ): @@ -927,7 +934,7 @@ def ep_dispatch_fwd( ) -@compute_on(_EP_COLLECTIVE_STREAM) +@_on_collective_stream def ep_combine_fwd( cfg: EpLayerConfig, handle_mem, expert_out, num_local_tokens, out_partition_spec=None ): @@ -943,7 +950,7 @@ def ep_combine_fwd( ) -@compute_on(_EP_COLLECTIVE_STREAM) +@_on_collective_stream def ep_dispatch_bwd( cfg: EpLayerConfig, handle_mem, @@ -965,7 +972,7 @@ def ep_dispatch_bwd( ) -@compute_on(_EP_COLLECTIVE_STREAM) +@_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( diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py index e4619d8670..8765a4fce0 100644 --- a/transformer_engine/jax/version_utils.py +++ b/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", From 9371934dc9c6865c158b6fb5c1319bb76a22a21b Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 22 Jul 2026 10:04:09 -0700 Subject: [PATCH 4/4] [JAX] Silence pylint not-callable on compute_on EP annotation Signed-off-by: Phuong Nguyen --- transformer_engine/jax/cpp_extensions/ep.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 0723db6ebe..806e7ae480 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -34,7 +34,7 @@ def _on_collective_stream(func): return func from jax.experimental.compute_on import compute_on - return compute_on("gpu_stream:collective")(func) + return compute_on("gpu_stream:collective")(func) # pylint: disable=not-callable __all__ = [