[JAX] Expert Parallelism: JAX primitives + VJPs - #3036

Merged

phu0ngng merged 24 commits into

NVIDIA:mainfrom

phu0ngng:phuong/ep-3-jax

Jun 27, 2026

Commits on Jun 24, 2026

jax: add EP bindings on pointer-keyed cache with EpLayerConfig and bf16 max_token_dtype
phu0ngng
committed
3ee9599
jax/ep: drop topk_weights from ep_combine; caller must pre-multiply
phu0ngng
committed
fb3300f
tests/jax/ep: mask uninitialized recv_tokens tail in dispatch_vjp
phu0ngng
committed
fa907ca
examples/jax/ep: add ep_bench.py + run_ep_bench.sh
phu0ngng
committed
dedbf86
examples/jax/ep: ep_moe.py runs --iters fwd+bwd steps (default 3)
phu0ngng
committed
453b399
jax/ep: tighten sharding contract, drop helpers, route bwd through TE with_sharding_constraint
phu0ngng
committed
4ee095c
jax/ep: derive ep_size and num_ep_groups from active mesh in ep_bootstrap
phu0ngng
committed
761d557
examples/jax/ep: rename ep_handle to layer_cfg in ep_moe.py (matches EpLayerConfig type)
phu0ngng
committed
82c9063
jax/ep: add primitive docstrings and silence missing-kwoa false positives (lint 10.00)
phu0ngng
committed
a314a45
jax/ep: apply black formatting (pre-commit hook output)
phu0ngng
committed
c0b280d
build_tools/jax: gate NCCL EP on NVTE_BUILD_WITH_NCCL_EP (default on); define NVTE_WITH_NCCL_EP
phu0ngng
committed
348cc36
jax/ep: collapse 5 FFI attr structs into single EpConfig
phu0ngng
committed
bf811b8
jax/ep: dedup _ep_outer_axis, normalize _ep_spec_ok, unify outer_abstract, drop dead helpers
phu0ngng
committed
c030618
jax/ep: apply clang-format and silence pylint unused-arg in lowering
phu0ngng
committed
e41f588
qa: wire NCCL EP tests into L1 (multi-process unittest) and L0 (ep_moe example) jax distributed suites
phu0ngng
committed
572889c
jax/ep: resolve test path from launcher SCRIPT_DIR and export LD_LIBRARY_PATH for libnccl_ep.so
phu0ngng
committed
943a970
[JAX] Generalize EP MoE example to N stacked layers via --num-layers
phu0ngng
committed
8818438
Unify NCCL EP build flag: rename to NVTE_WITH_NCCL_EP, share arch check via nccl_ep_enabled()
phu0ngng
committed
a3378db
Add arg descriptions to ep_bootstrap docstring
phu0ngng
committed
06440a6
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
phu0ngng
committed
ebf5d96
L0_jax_unittest: exclude multi-process EP test on single-node runner
phu0ngng
committed
9df769a

Commits on Jun 25, 2026

Guard against None out_partition_spec in EP combine/dispatch-bwd partition methods
phu0ngng
committed
db8d43b

Commits on Jun 26, 2026

Fix test_two_layer_dispatch_no_handle_aliasing: add A->B data edge
phu0ngng
committed
0a29355

Commits on Jun 27, 2026

Merge branch 'main' into phuong/ep-3-jax
phu0ngng
authored
9187d81