TE EP integration to MoEBlock - #3116

Merged

tdophung merged 23 commits into

NVIDIA:mainfrom

tdophung:teddy/te_ep_integration

Jul 10, 2026

Commits on Jul 7, 2026

[JAX] Resync onto upstream PR #3036, restore TE-EP-only MoE block
tdophung
committed
cba9717
tests/jax: trim TE-EP MoE suite (drop bootstrap, flax-wrapper, bias-zero)
tdophung
committed
006902e
jax/router: fix two bwd custom_partitioning bugs (aux-loss rank, topk closure)
tdophung
committed
51a046a
jax/ep: skip size-1 dp/fsdp axis in _ep_outer_axis
tdophung
committed
375d080
jax/flax: realign _MoEBlock with post-resync moe() signature
tdophung
committed
f3f77f7
jax/moe: plumb token_counts to grouped_gemm and zero 0-token wgrad slices
tdophung
committed
49058ff
jax/flax,tests: rename use_bias/use_expert_bias for symmetry (PR #3116)
tdophung
committed
3e77957
jax/moe: address PR #3116 review feedback (hardcode align + expand inline justifications)
tdophung
committed
2f54d83
jax/moe: strip PR-response framing from comments; drop sparse_probs NaN sanitizer
tdophung
committed
09d5f78
jax/moe: drop fp32 island around silu+multiply (fwd, bwd, reference)
tdophung
committed
9d959f2
remove useless comments
tdophung
committed
b7d3a85
tests/jax: remove legacy MoE VJP tests + launcher; point CI at TE-EP successor
tdophung
committed
ff50f44
jax/moe: swap _Ctx to @flax.struct.dataclass, drop manual pytree boilerplate
tdophung
committed
641b6b8
jax/moe: drop bwd recv_topk_weights NaN sanitizer; trust the dispatch contract
tdophung
committed
055b72c
jax/moe: assert output dtype; tests cover d_x parity (dtype + values)
tdophung
committed
e670511
tests/jax/test_te_ep_moe: strip docstring to just "what this suite covers"
tdophung
committed
518e177
jax/moe: address TE EP alignment review feedback
tdophung
committed
bd052ce
jax/moe: fix early topk weighting padded-slot masking
tdophung
committed
c2869ed
jax/moe: remove unused EP mesh size
tdophung
committed
2dd0c6c

Navigation Menu

Uh oh!

TE EP integration to MoEBlock - #3116

TE EP integration to MoEBlock#3116
tdophung merged 23 commits into
NVIDIA:mainfrom
tdophung:teddy/te_ep_integration

Commits on Jul 7, 2026

[JAX] Resync onto upstream PR #3036, restore TE-EP-only MoE block

tests/jax: trim TE-EP MoE suite (drop bootstrap, flax-wrapper, bias-zero)

jax/router: fix two bwd custom_partitioning bugs (aux-loss rank, topk closure)

jax/ep: skip size-1 dp/fsdp axis in _ep_outer_axis

jax/flax: realign _MoEBlock with post-resync moe() signature

jax/moe: plumb token_counts to grouped_gemm and zero 0-token wgrad slices

jax/flax,tests: rename use_bias/use_expert_bias for symmetry (PR #3116)

jax/moe: address PR #3116 review feedback (hardcode align + expand inline justifications)

jax/moe: strip PR-response framing from comments; drop sparse_probs NaN sanitizer

jax/moe: drop fp32 island around silu+multiply (fwd, bwd, reference)

remove useless comments

tests/jax: remove legacy MoE VJP tests + launcher; point CI at TE-EP successor

jax/moe: swap _Ctx to @flax.struct.dataclass, drop manual pytree boilerplate

jax/moe: drop bwd recv_topk_weights NaN sanitizer; trust the dispatch contract

jax/moe: assert output dtype; tests cover d_x parity (dtype + values)

tests/jax/test_te_ep_moe: strip docstring to just "what this suite covers"

jax/moe: address TE EP alignment review feedback

jax/moe: fix early topk weighting padded-slot masking

jax/moe: remove unused EP mesh size

Commits on Jul 8, 2026

jax/moe: tighten TE EP recv capacity bound

jax/moe: simplify late TE EP weighting

[pre-commit.ci] auto fixes from pre-commit.com hooks

jax/moe: reduce padded-slot recv weight masking