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

[phu0ngng]

Description

Pin the TE-JAX Expert Parallelism dispatch/combine ops to XLA's high-priority collective stream so the scheduler serializes them with native collectives instead of overlapping.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR pins the TE-JAX Expert Parallelism dispatch/combine ops to XLA's gpu_stream:collective stream so the XLA scheduler serializes them with native collectives rather than overlapping, fixing potential ordering issues in multi-GPU EP workloads.

• Adds is_collective_stream_supported() to version_utils.py (gated on JAX ≥ 0.10.0) and a _on_collective_stream decorator in ep.py that lazily imports jax.experimental.compute_on — avoiding fatal errors on older JAX at module load time.
• Applies the decorator to all five public EP helpers (ep_prepare, ep_dispatch_fwd, ep_combine_fwd, ep_dispatch_bwd, ep_combine_bwd), with the wrapping decision made once at module import time.
• Adds test_z_dispatch_combine_on_collective_stream which compiles the forward HLO and asserts every te_ep_* custom call carries _xla_stream_annotation=\"collective\"; the test is skipped automatically on JAX < 0.10.0.

Confidence Score: 5/5

Safe to merge — the change is fully additive, the JAX version gate ensures no breakage on older installs, and the new test validates the annotation end-to-end for the forward path.

The decorator is a clean, well-isolated addition that falls back to a no-op on older JAX. All five EP primitives are decorated, the version guard is consistent with existing patterns in version_utils.py, and the lazy import of compute_on correctly prevents import-time failures on JAX < 0.10.0.

The new HLO test in tests/jax/test_multi_process_ep.py only covers the forward pass, leaving ep_dispatch_bwd and ep_combine_bwd without annotation verification.

Important Files Changed

Filename	Overview
transformer_engine/jax/version_utils.py	Adds is_collective_stream_supported() and _COLLECTIVE_STREAM_MIN_JAX_VERSION = "0.10.0" following the established version-check pattern. Correctly delegates to the already-cached jax_version_meet_requirement.
transformer_engine/jax/cpp_extensions/ep.py	Introduces _on_collective_stream decorator applied to all five public EP helpers. Decorator defers compute_on import so older JAX versions don't fail at load time.
tests/jax/test_multi_process_ep.py	New test compiles forward HLO and asserts every te_ep_* custom call carries _xla_stream_annotation="collective". Correctly gated behind is_collective_stream_supported().

_{Reviews (3): Last reviewed commit: "[JAX] Gate EP collective-stream annotati..." | Re-trigger Greptile}

[phu0ngng]

/te-ci JAX L1

[jberchtold-nvidia]

What JAX version was this added in? Should we guard with something like

try:
   from jax.experimental.compute_on import compute_on
except ImportError:
   def compute_on(stream_name: str):
       """ compute_on shim for older JAX versions """
       # Make @compute_on annotation do nothing to the annotated function
       return lambda f: f

[phu0ngng]

Good suggestion for versioning.

The import compute_on is available since 2024, I think, so this check won't help here.
I added a guard for JAX 0.10.0.

[phu0ngng]

/te-ci JAX L1