[JAX] Migrate EP collective-stream annotation to compute_on2

tests/jax/test_multi_process_ep.py

@@ -19,6 +19,7 @@
 """
 
 import os
+import re
 import sys
 import unittest
 
@@ -673,9 +674,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_on2
+        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)
@@ -703,18 +704,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):

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
@@ -32,9 +33,22 @@ def _on_collective_stream(func):
     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
+    from jax.experimental.compute_on import compute_on2
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        # compute_on2 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_on2(
+            compute_type="gpu_stream:collective",
+            out_memory_spaces=jax.memory.Space.Device,
+        )(lambda: func(*args, **kwargs))
+        return annotated()
+
+    return wrapper
 
 
 __all__ = [