[JAX] Migrate EP collective-stream annotation to region-based compute_on

tests/jax/test_multi_process_ep.py

@@ -22,6 +22,7 @@
 """
 
 import os
+import re
 import sys
 import unittest
 
@@ -689,9 +690,9 @@ def run(idx, toks, w):
         "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."""
+        """Every EP FFI custom call must run on the collective stream. compute_on
+        puts the annotation on the async wrapper XLA generates, so assert each EP
+        call is reachable from a wrapper that carries it."""
         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)
@@ -719,18 +720,50 @@ def run(idx, toks, w):
 
             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}")
+        # Parse the HLO into computations and follow the call graph: a call
+        # carrying the collective-stream annotation places its callee (and every
+        # nested callee) on the collective stream.
+        comps = {}
+        cur = None
+        for line in hlo.splitlines():
+            stripped = line.strip()
+            header = re.match(r"(?:ENTRY\s+)?(%[\w.\-]+)\s*\(", stripped)
+            if header and stripped.endswith("{"):
+                cur = header.group(1)
+                comps[cur] = []
+            elif stripped == "}":
+                cur = None
+            elif cur is not None:
+                comps[cur].append(line)
+
+        callees = lambda l: re.findall(r"(?:calls|to_apply)=(%[\w.\-]+)", l)
+        edges = {c: {x for l in ls for x in callees(l)} for c, ls in comps.items()}
+
+        collective = set()
+        for ls in comps.values():
+            for l in ls:
+                if '_xla_stream_annotation="collective"' in l.replace(" ", ""):
+                    collective.update(callees(l))
+        stack = list(collective)
+        while stack:
+            for callee in edges.get(stack.pop(), ()):
+                if callee not in collective:
+                    collective.add(callee)
+                    stack.append(callee)
+
+        ep_calls = [
+            (c, l) for c, ls in comps.items() for l in ls if 'custom_call_target="te_ep_' in l
+        ]
+        self.assertTrue(ep_calls, 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(" ", "")
+            for c, l in ep_calls
+            if c not in collective
+            and '_xla_stream_annotation="collective"' not in l.replace(" ", "")
         ]
         self.assertFalse(
             missing,
-            "te_ep_* custom calls missing collective-stream annotation:\n" + "\n".join(missing),
+            "te_ep_* custom calls not on the collective stream:\n" + "\n".join(missing),
         )
 
     def test_z_no_unexpected_reshard_in_hlo_bwd(self):

tests/jax/test_te_ep_moe.py

@@ -140,7 +140,7 @@ def _read_mp_options():
     ("exp", EP_AXIS),
     ("embed", FSDP_AXIS),
     ("mlp", None),
-    ("batch", (EP_AXIS, FSDP_AXIS)),
+    ("batch", (FSDP_AXIS, EP_AXIS)),
 )
 
 # Small shapes so the parity tests stay tight on bf16. The block still
@@ -364,6 +364,7 @@ def _make_block(
     use_expert_routing_bias=False,
     score_function="softmax",
     expert_bias_init=None,
+    input_axes=("batch", None, None),
 ):
     kwargs = dict(
         num_experts=NUM_EXPERTS,
@@ -375,6 +376,7 @@ def _make_block(
         use_expert_routing_bias=use_expert_routing_bias,
         score_function=score_function,
         dtype=DTYPE,
+        input_axes=input_axes,
     )
     # Custom expert_bias_init lets tests inject a non-zero expert_bias without
     # poking variables['params'] post-init.

transformer_engine/jax/cpp_extensions/ep.py

@@ -14,6 +14,7 @@
     compound ``(dp, ep)`` axis on the leading dim.
 """
 
+import functools
 from dataclasses import dataclass
 
 import jax
@@ -34,7 +35,20 @@ def _on_collective_stream(func):
         return func
     from jax.experimental.compute_on import compute_on
 
-    return compute_on("gpu_stream:collective")(func)  # pylint: disable=not-callable
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        # compute_on traces its callee and abstract-evals every argument, so it
+        # cannot take the static EpLayerConfig/PartitionSpec args directly. Wrap
+        # a nullary thunk that closes over them; the array operands are captured
+        # as consts and lifted to real operands, outputs stay on device. XLA
+        # async-wraps the resulting call onto the collective stream.
+        annotated = compute_on(
+            compute_type="gpu_stream:collective",
+            out_memory_spaces=jax.memory.Space.Device,
+        )(lambda: func(*args, **kwargs))
+        return annotated()
+
+    return wrapper
 
 
 __all__ = [

transformer_engine/jax/moe.py

@@ -751,6 +751,8 @@ def _body(*args):
             num_local_tokens=(B, S),
             out_partition_spec=out_partition_spec,
         )
+    # output of MLP should be sharded the same way as the activation input
+    output = with_sharding_constraint_by_logical_axes(output, input_axes)
 
     (
         casted_sorted_x_lhs_trans,