[JAX] Keep the routing map format alive and EP multiprocess tests in L2 Jax dist

[KshitijLakhani]

Description

Keeps the routing map alive so that the nested function can use it

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

Please list the changes introduced in this PR:

• Change A
• Change B

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

[KshitijLakhani]

/te-ci jax L2

[greptile-apps]

Greptile Summary

This PR fixes a NameError bug in the JAX MoE backward pass sharding logic where routing_map_format was incorrectly deleted before a nested closure could capture it, and adds the NCCL EP multi-process test suite to the distributed CI tier while excluding it from the single-process tier.

• router.py: Removes routing_map_format from the del statement in FusedTopkWithScoreFunctionBwdPrimitive.partition. Without this fix, sharded_impl captures the deleted name and raises NameError at call time when distributed sharding is active.
• test.sh (distributed): Adds multi_process_launch_ep.sh invocation, which self-skips when fewer than 4 GPUs or NVLink is absent.
• test.sh (unit): Adds --ignore flag so test_multi_process_ep.py is not collected by the single-process pytest run.

Confidence Score: 5/5

Safe to merge — the one-line change in router.py correctly restores a variable that was inadvertently deleted before its nested closure could use it.

The router.py change removes a name from a del statement that was actively breaking the sharded_impl closure in the MoE backward pass sharding path. The fix is minimal, targeted, and does not affect any other code paths. The test.sh changes are additive and internally consistent with the existing CI structure.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/jax/cpp_extensions/router.py	Single-line fix: removes routing_map_format from the del statement so the nested sharded_impl closure can capture it; change is correct and targeted.
qa/L2_jax_distributed_unittest/test.sh	Adds NCCL EP multi-process launcher to the distributed test suite; set -xe propagates failures correctly.
qa/L2_jax_unittest/test.sh	Adds --ignore for test_multi_process_ep.py so multi-process tests are not run in the single-GPU pytest tier.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["partition(topk, ..., routing_map_format, mesh, arg_infos, result_infos)"]
    B["del result_infos\n(routing_map_format kept alive ✓)"]
    C["Compute out_sharding & arg_shardings"]
    D["def sharded_impl(routing_map, intermediate, grad_probs):\n  calls impl(..., routing_map_format)"]
    E["return mesh, sharded_impl, out_sharding, arg_shardings"]
    F["JAX calls sharded_impl at dispatch time"]
    G["FusedTopkWithScoreFunctionBwdPrimitive.impl(..., routing_map_format)"]

    A --> B
    B --> C
    C --> D
    D --> E
    E --> F
    F --> G

    style B fill:#d4edda,stroke:#28a745
    style D fill:#fff3cd,stroke:#ffc107

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A["partition(topk, ..., routing_map_format, mesh, arg_infos, result_infos)"]
    B["del result_infos\n(routing_map_format kept alive ✓)"]
    C["Compute out_sharding & arg_shardings"]
    D["def sharded_impl(routing_map, intermediate, grad_probs):\n  calls impl(..., routing_map_format)"]
    E["return mesh, sharded_impl, out_sharding, arg_shardings"]
    F["JAX calls sharded_impl at dispatch time"]
    G["FusedTopkWithScoreFunctionBwdPrimitive.impl(..., routing_map_format)"]

    A --> B
    B --> C
    C --> D
    D --> E
    E --> F
    F --> G

    style B fill:#d4edda,stroke:#28a745
    style D fill:#fff3cd,stroke:#ffc107

Loading

_{Reviews (2): Last reviewed commit: "Fix incorrectly launched multi process E..." | Re-trigger Greptile}

[KshitijLakhani]

/te-ci jax L2

[jberchtold-nvidia]

LGTM, thanks for making this fix Kshitij!

[KshitijLakhani]

CI though not fully complete looks good to merge
Pipeline: 352279726