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

[phu0ngng]

Description

Support single-process, multi-device JAX EP using XLA-borrowed NCCL communicators.

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

Changes

• Manage EP initialization, state, and teardown per device.
• Size-handle buffers across local devices and reject undersized buffers.
• Update bootstrap, launchers, examples, and tests for multi-device execution.

Testing

• EP guard/grouping tests: 7 passed, 1 skipped.
• All pre-commit checks passed.

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]

Confidence Score: 5/5

The PR appears safe to merge; no actionable correctness, security, or repository-rule violations remain.

Summary

This PR adds JAX expert-parallel support for multiple devices controlled by one process using XLA-borrowed NCCL communicators.

• Stores and tears down EP backend state per CUDA device.
• Bootstraps borrowed communicators according to JAX mesh EP replica groups.
• Validates handle-buffer sizes before NCCL EP operations.
• Adds single-process examples, multi-device launch configurations, and EP lifecycle/grouping tests.
• Preserves the positional ep_bootstrap API highlighted in the previous review.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[JAX ep_bootstrap] --> B{Multiple local devices?}
    B -->|No| C[Self-hosted or borrowed communicator]
    B -->|Yes| D{XLA borrowed communicator available?}
    D -->|No| E[Reject bootstrap]
    D -->|Yes| F[Derive EP replica groups from mesh]
    F --> G[XLA bootstrap executes on each device]
    G --> H[Initialize per-device EPBackend]
    H --> I[Query maximum handle size across devices]
    I --> J[Prepare / dispatch / combine]
    J --> K[Finalize all device backends concurrently]

Loading

_{Reviews (4) · Last reviewed commit: "[JAX] Preserve positional EP bootstrap a..."}

[jberchtold-nvidia]

I like this mocking of tests to test our error paths too!

[jberchtold-nvidia]

Looks great! Thanks!

Once CI passes, all good to merge from my side

[phu0ngng]

/te-ci L1 JAX