[JAX] Support single-process-multi-device EP via XLA-borrowed communicator - #3522

Open

phu0ngng wants to merge 3 commits into

NVIDIA:mainfrom

phu0ngng:ep_jax_single_process_backup

Commits on Sep 15, 2026

[JAX] Support single-process, multi-device NCCL EP via borrowed comm
phu0ngng
committed
ad18bb1
Update transformer_engine/jax/ep.py

phu0ngng
and
greptile-apps[bot]
authored
cb040f9
[JAX] Preserve positional EP bootstrap arguments
phu0ngng
committed
26bba3f