[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, with no outstanding correctness, security, or repository-rule findings.

Summary

This PR adds single-process, multi-device JAX expert parallelism using XLA-borrowed NCCL communicators.

• Reworks the native EP backend from process-global state to per-CUDA-device state.
• Initializes borrowed communicators independently for each local device and coordinates teardown.
• Sizes routing-handle buffers across initialized devices and rejects undersized buffers.
• Updates JAX bootstrap validation, examples, launchers, and distributed test configurations.
• Preserves the existing positional ep_bootstrap calling convention addressed by the resolved previous review thread.

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 required]
    D --> E[FFI obtains communicator per device]
    E --> F[Initialize per-device EPBackend]
    F --> G[Compute maximum handle size]
    G --> H[Prepare / dispatch / combine per device]
    H --> I[ep_finalize]
    I --> J[Clear executable caches]
    J --> K[Shut down all local backends]

Loading

_{Reviews (3) · 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