[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 all five TE-JAX Expert Parallelism FFI helpers (ep_prepare, ep_dispatch_fwd/bwd, ep_combine_fwd/bwd) to XLA's high-priority collective stream via a new _on_collective_stream decorator, preventing the XLA scheduler from interleaving them with native NCCL collectives. A version gate (is_collective_stream_supported, JAX >= 0.10.0) is added to version_utils.py so older JAX versions see no change.

• transformer_engine/jax/cpp_extensions/ep.py — introduces _on_collective_stream, applied at module-load time; lazily imports jax.experimental.compute_on.compute_on only when supported.
• transformer_engine/jax/version_utils.py — adds _COLLECTIVE_STREAM_MIN_JAX_VERSION = \"0.10.0\" and is_collective_stream_supported(), exported in __all__.
• tests/jax/test_multi_process_ep.py — adds a forward-only HLO test that compiles a dispatch/combine trace and asserts every te_ep_* custom call carries _xla_stream_annotation=\"collective\"; test is skipped on unsupported JAX.

Confidence Score: 5/5

Safe to merge — the feature is gated behind a version check so older JAX installations see no behavioural change, and the new test validates the annotation on the forward path.

The change is an opt-in scheduling hint applied at module import time. The version guard correctly isolates older runtimes, the decorator mechanism is straightforward, and the new HLO test confirms the annotation appears in compiled output. No logic errors or API misuses were found.

The ep_dispatch_bwd and ep_combine_bwd backward ops are decorated but their annotation is not verified by any test — worth adding a backward-path assertion to test_z_no_unexpected_reshard_in_hlo_bwd or a companion test.

Important Files Changed

Filename	Overview
transformer_engine/jax/version_utils.py	Adds is_collective_stream_supported() gating gpu_stream:collective on JAX >= 0.10.0, and exports it in __all__. Consistent with existing patterns; delegates to the cached jax_version_meet_requirement.
transformer_engine/jax/cpp_extensions/ep.py	Introduces _on_collective_stream decorator applied at module-load time to all five public EP helpers; lazily imports compute_on only on supported JAX versions. Backward ops (ep_dispatch_bwd, ep_combine_bwd) are decorated but the annotation isn't verified by the new test.
tests/jax/test_multi_process_ep.py	Adds test_z_dispatch_combine_on_collective_stream, guarded by the version check, that compiles a forward-only jit trace and asserts every te_ep_* custom call carries _xla_stream_annotation="collective" in the emitted HLO.

Sequence Diagram

sequenceDiagram
    participant TE as TE ep.py import
    participant VU as version_utils
    participant XLA as XLA scheduler

    TE->>VU: "is_collective_stream_supported()?"
    VU-->>TE: "JAX >= 0.10.0 → True"
    TE->>TE: "compute_on(gpu_stream:collective)(func) wraps all 5 EP helpers"

    Note over TE,XLA: At JIT compile time

    TE->>XLA: "EpDispatchPrimitive.bind() with _xla_stream_annotation=collective"
    XLA-->>XLA: "schedule on NCCL collective stream"
    TE->>XLA: "EpCombinePrimitive.bind() with _xla_stream_annotation=collective"
    XLA-->>XLA: "serialised with native all-reduce / all-to-all ops"

Loading

_{Reviews (5): Last reviewed commit: "Merge branch 'main' into ep/jax_col_stre..." | Re-trigger Greptile}

[phu0ngng]

/te-ci JAX L1

[phu0ngng]

/te-ci JAX L1

[jberchtold-nvidia]

LGTM pending CI, thanks!