[Regression] `CollectivePipeliner` GPU pass present in JAX 0.4.38 is missing from JAX 0.10.1 — pipelining flags accepted but produce no pass stages; class symbols absent from binaries

## Summary The XLA collective pipeliner pass (`xla/service/collective_pipeliner.cc`) was present and runnable in the GPU plugin shipped with `jax==0.4.38` (confirmed via `strings`/`nm` of `jax_plugins/xla_cuda12/xla_cuda_plugin.so` — see "Cross-check on older JAX" below). In `jax==0.10.1` / `jaxlib==0.10.1` (current stable as of June 2026), the pass is no longer invoked on GPU regardless of which `--xla_gpu_enable_pipelined_*` flag is set: - `--xla_gpu_enable_pipelined_all_reduce=true` (and the `_all_gather` / `_reduce_scatter` variants) are accepted by the flag parser without error - HLO pass dumps (`XLA_FLAGS=--xla_dump_hlo_pass_re=.*`) contain **no `collective-pipeliner-forward` or `-backward` stages** in any compiled module - `nm -D` / `strings` on every `.so` shipped in `site-packages/jaxlib/` and `site-packages/jax_plugins/` shows **no `CollectivePipeliner` symbols** — the class is not linked into the binary - The umbrella flag `--xla_gpu_enable_pipelined_collectives=true` is rejected as "Unknown flag" by the parser - `--xla_gpu_run_post_layout_collective_pipeliner=true` is also rejected as "Unknown flag" (though it appears as a proto field name string in `libjax_common.so`) The pass is still referenced in XLA source (`xla/service/collective_pipeliner.cc`), so it exists upstream — but appears not built into the GPU plugin shipped with current JAX stable. The "Cross-check on older JAX" section below shows the same wheels installed via the same `pip install "jax[cuda12]==<version>"` command have the symbols in 0.4.38 but not in 0.10.1, narrowing the regression to a build-config change between those two releases. ## Environment | Component | Version | |---|---| | OS | CentOS Stream 9, Linux 6.13.2 x86_64 | | GPU | NVIDIA H100, driver 580.82.07 | | Python | 3.11.15 | | `jax` | 0.10.1 | | `jaxlib` | 0.10.1 | | `jax-cuda12-plugin` | 0.10.1 | | `jax-cuda12-pjrt` | 0.10.1 | | NCCL | 2.30.7 + CUDA 12.9 | | Hardware | 8× H100 single node | ## Reproduction ### Step 0 — fresh conda env (no environment contamination) ```bash conda create -n xla_repro_clean python=3.11 -y ~/.conda/envs/xla_repro_clean/bin/pip install --upgrade pip ~/.conda/envs/xla_repro_clean/bin/pip install "jax[cuda12]==0.10.1" ``` All subsequent commands use `~/.conda/envs/xla_repro_clean/bin/python` to guarantee a clean dependency stack. ### Step 1 — minimal repro script Save as `/tmp/repro.py` — Megatron-style TP-MLP scan that produces a `while` loop containing an `all-reduce`: ```python import jax import jax.numpy as jnp import numpy as np from jax.sharding import Mesh, PartitionSpec as P, NamedSharding from jax import lax assert jax.device_count() >= 8, f"need 8 devices, got {jax.device_count()}" devices = jax.devices()[:8] mesh = Mesh(np.array(devices).reshape(8), ("tp",)) D, D_FF, N_LAYERS, B, S = 4096, 16384, 12, 4, 128 DTYPE = jnp.bfloat16 x = jnp.ones((B, S, D), dtype=DTYPE) w_ups = jnp.ones((N_LAYERS, D, D_FF), dtype=DTYPE) w_downs = jnp.ones((N_LAYERS, D_FF, D), dtype=DTYPE) x_sh = NamedSharding(mesh, P(None, None, None)) wu_sh = NamedSharding(mesh, P(None, None, "tp")) # D_FF sharded on output wd_sh = NamedSharding(mesh, P(None, "tp", None)) # D_FF sharded on contracting x = jax.device_put(x, x_sh) w_ups = jax.device_put(w_ups, wu_sh) w_downs = jax.device_put(w_downs, wd_sh) def scan_body(x, ws): w_up, w_down = ws h = jax.nn.gelu(x @ w_up) @ w_down # produces all-reduce on contracting dim return h, None def model(x, w_ups, w_downs): y, _ = lax.scan(scan_body, x, (w_ups, w_downs)) return y f = jax.jit(model, in_shardings=(x_sh, wu_sh, wd_sh), out_shardings=x_sh) hlo = f.lower(x, w_ups, w_downs).compile().as_text() print(f"all-reduce: {hlo.count('all-reduce')}, while: {hlo.count('while(')}") y = f(x, w_ups, w_downs).block_until_ready() print("Done.", y.shape) ``` ### Step 2 — verify the compiled HLO has the right structure (`while` + `all-reduce`) ```bash ~/.conda/envs/xla_repro_clean/bin/python /tmp/repro.py ``` Output: ``` all-reduce: 8, while: 1 Done. (4, 128, 4096) ``` ### Step 3 — verify no pipeliner pass stages are dumped ```bash rm -rf /tmp/xla_dump && mkdir /tmp/xla_dump XLA_FLAGS="--xla_dump_to=/tmp/xla_dump --xla_dump_hlo_pass_re=.* \ --xla_gpu_enable_pipelined_all_reduce=true \ --xla_gpu_enable_pipelined_all_gather=true \ --xla_gpu_enable_pipelined_reduce_scatter=true" \ ~/.conda/envs/xla_repro_clean/bin/python /tmp/repro.py ls /tmp/xla_dump/ | grep -E 'collective-pipeliner-(forward|backward)' | wc -l # Expected: > 0 ; Actual: 0 ``` ### Step 4 — verify the `CollectivePipeliner` class is not linked into any `.so` ```bash find ~/.conda/envs/xla_repro_clean/lib/python3.11/site-packages -name "*.so" 2>/dev/null \ | while read f; do nm -D "$f" 2>/dev/null | grep -q CollectivePipeliner && echo "$f" done # Expected: at least one .so listed ; Actual: no output ``` ### Step 5 — verify the umbrella + post-layout flags are unrecognized ```bash XLA_FLAGS="--xla_gpu_enable_pipelined_collectives=true" \ ~/.conda/envs/xla_repro_clean/bin/python /tmp/repro.py 2>&1 | head -1 # Output: # F0612 08:30:24.243982 ...parse_flags_from_env.cc:234] Unknown flag in XLA_FLAGS: --xla_gpu_enable_pipelined_collectives=true XLA_FLAGS="--xla_gpu_run_post_layout_collective_pipeliner=true" \ ~/.conda/envs/xla_repro_clean/bin/python /tmp/repro.py 2>&1 | head -1 # Output: # F0612 08:36:33.970901 ...parse_flags_from_env.cc:234] Unknown flag in XLA_FLAGS: --xla_gpu_run_post_layout_collective_pipeliner=true ``` ### Single-GPU / no-GPU verification path The headline finding (pass missing from binary) can be verified by anyone without GPU access: ```bash conda create -n xla_repro_nogpu python=3.11 -y ~/.conda/envs/xla_repro_nogpu/bin/pip install "jax[cuda12]==0.10.1" SO=$(~/.conda/envs/xla_repro_nogpu/bin/python -c \ "import jaxlib, os; print(os.path.join(os.path.dirname(jaxlib.__file__), 'libjax_common.so'))") nm -D "$SO" | grep CollectivePipeliner # Expected: at least one symbol ; Actual: no output ``` (The CUDA12 wheel installs cleanly on hosts without GPUs; only Steps 2-3 require actual GPUs.) ## Expected vs Actual | Behavior | Expected | Actual | |---|---|---| | `collective-pipeliner-forward` stage appears in pass dumps | yes | no | | `collective-pipeliner-backward` stage appears in pass dumps | yes | no | | `CollectivePipeliner` class linked into binary | yes | no | | `--xla_gpu_enable_pipelined_collectives=true` accepted | yes (umbrella in newer XLA) | "Unknown flag" | | `--xla_gpu_enable_pipelined_all_reduce=true` triggers pipeliner | yes | accepted, no effect | ## Cross-check on older JAX Installed `jax==0.4.38` in a separate fresh env to verify the pass is supposed to exist. Confirmed `CollectivePipeliner` symbols **are** present in `jax_plugins/xla_cuda12/xla_cuda_plugin.so` in that version: ``` $ strings xla_cuda_plugin.so | grep -E "collective-pipeliner|CollectivePipeliner" collective-pipeliner-backward collective-pipeliner-forward collective-pipeliner-forwardsink N3xla19CollectivePipelinerE external/xla/xla/service/collective_pipeliner.cc ``` So the pass exists upstream and was present in earlier JAX versions. It appears to have been dropped from the build between 0.4.38 and 0.10.1. ## Questions 1. Is the `CollectivePipeliner` pass intentionally removed from the JAX-bundled XLA GPU build in 0.10.1? If so: - What is the recommended path for cross-iteration collective overlap on GPU? - Has the functionality been consolidated into another pass (e.g. latency-hiding-scheduler)? If so, please document. - Should the dead flags (`--xla_gpu_enable_pipelined_all_reduce` etc.) be removed or upgraded to errors to avoid confusion? 2. If unintentional, what build flag is needed to re-include the pass? `--xla_gpu_run_post_layout_collective_pipeliner` appears to be a proto field but not a CLI flag in this build. 3. Is the recommended way to get pipelining on GPU now via shardy/`shard_map` + manual primitives, or some other mechanism?