Multi-stream windowed einsum fails HLO verifier when a per-device shard > ~1024 — REGRESSION new in jax/jaxlib 0.10.2 (CUDA)

### Description ## Version bisect (this is a 0.10.2 regression) Same minimal repro, same 8×H100, per-device N shard = 2048 (`repro.py 2048`): | jax / jaxlib | per-dev 1024 | per-dev 2048 (>1024) | |---|---|---| | 0.10.0 | compiles | **compiles** | | 0.10.1 | compiles | **compiles** | | 0.10.2 | compiles | **FAILED_PRECONDITION (regression)** | | 0.8.0 | compiles | SIGSEGV (older, different failure) | So the verifier failure was **introduced between jaxlib 0.10.1 (2026-05-20) and 0.10.2 (2026-06-17)** — a bisectable window. On 0.10.1 the large-per-device windowed programs compile, run correctly, and multi-stream overlap engages (it beats the single-stream control by ~8% at large N). **Flags:** `--xla_gpu_multi_streamed_windowed_einsum=true --xla_gpu_threshold_for_windowed_einsum_mib=0` ## Summary With multi-stream windowed einsum forced, XLA's GPU backend fails compilation: ``` jax.errors.JaxRuntimeError: FAILED_PRECONDITION: during context [hlo verifier]: Asynchronous computation async_computation expected to contain only the root and parameter instructions. Failed after pipeline-start ``` whenever a **per-device sharded matmul dimension exceeds ~1024** — i.e. `M/tp > ~1024` or `N/tp > ~1024` for a tensor-parallel `all-gather(x[M,K]) → x_full @ w1[K, N/tp]`. It reproduces with a single windowed collective-matmul (one dot + a `tanh` consumer); it is not specific to chained loops. ## Trigger boundary (single-layer, `tanh(x @ w1)`, tp=8 unless noted) | shape | M/tp | N/tp | result | |---|---|---|---| | M=8192, K=8192, N=8192 | 1024 | 1024 | compiles | | M=8192, K=8192, N=16384 | 1024 | **2048** | **FAILED_PRECONDITION** | | M=16384, K=8192, N=8192 | **2048** | 1024 | **FAILED_PRECONDITION** | | M=16384, K=16384, N=16384 | 2048 | 2048 | **FAILED_PRECONDITION** | | M=4096, N=4096/8192 | ≤512 | ≤1024 | compiles | | M=8192, N=8192, K ∈ {4096,16384,32768,65536} | 1024 | 1024 | compiles (K irrelevant) | | **M=8192,K=8192,N=8192 @ tp=4** | **2048** | **2048** | **FAILED_PRECONDITION** | | **M=8192,K=8192,N=8192 @ tp=2** | **4096** | **4096** | **FAILED_PRECONDITION** | The tp rows are decisive: the identical 8192³ program compiles at tp=8 (per-device 1024) but fails at tp=4/tp=2 (per-device 2048/4096) — so the trigger is the **per-device shard size**, not absolute dimensions. Threshold is in (1024, 2048]. **K does not matter.** The 2-layer chained MLP `(x@w1)@w2` fails at **all** shapes. ## Minimal repro ```python import os, sys os.environ["XLA_FLAGS"] = ("--xla_gpu_multi_streamed_windowed_einsum=true " "--xla_gpu_threshold_for_windowed_einsum_mib=0") os.environ.setdefault("JAX_PLATFORMS", "cuda") import jax, jax.numpy as jnp from jax.sharding import Mesh, NamedSharding, PartitionSpec as P tp = len(jax.devices()); assert tp >= 2 per_dev = int(sys.argv[1]) if len(sys.argv) > 1 else 2048 # >1024 fails, <=1024 ok M, K, N = 1024 * tp, 8192, per_dev * tp mesh = Mesh(jax.devices(), ("tp",)) k1, k2 = jax.random.split(jax.random.PRNGKey(0), 2) x = jax.device_put(jax.random.normal(k1, (M, K), jnp.bfloat16), NamedSharding(mesh, P("tp", None))) w1 = jax.device_put(jax.random.normal(k2, (K, N), jnp.bfloat16), NamedSharding(mesh, P(None, "tp"))) hs = NamedSharding(mesh, P(None, "tp")) def f(x, w1): return jnp.tanh(jax.lax.with_sharding_constraint(x @ w1, hs)) jax.jit(f, out_shardings=hs)(x, w1).block_until_ready() ``` ```bash python3 -m venv ~/v && source ~/v/bin/activate && pip install "jax[cuda12]==0.10.2" JAX_TRACEBACK_FILTERING=off python repro.py 2048 # per-device shard 2048 -> FAILED_PRECONDITION python repro.py 1024 # per-device shard 1024 -> OK ``` Verified: `2048 → FAILED_PRECONDITION (Failed after pipeline-start)`; `1024 → OK: compiled + ran`. ## Failing stage & history - Fails in `backend_compile_and_load` → HLO verifier, after the `pipeline-start` pass (multi-stream windowed-einsum pipelining). The pass emits an `async_computation` containing more than the root+parameter instructions the verifier allows. - **0.10.0 / 0.10.1: no failure** — same programs compile at per-device shards well above 1024. So this is a regression in 0.10.2, not a long-standing limit. - On the much older **jaxlib 0.8.0** these programs **SIGSEGV'd** in `backend_compile_and_load` (a separate, older failure) — unrelated to the 0.10.1→0.10.2 regression. ## Impact Blocks compiler-driven windowed collective-matmul overlap at any realistic per-device tile (>1024) on 0.10.2. The only compilable shapes there (per-device ≤1024) are too small to hide the all-gather under compute. On 0.10.1 (no bug) the compute-heavy shapes run and overlap works (multi-stream beats the non-overlapped control by ~8% at large N), though on GPU/NVLink it still doesn't beat the monolithic baseline for the shapes tested — a separate perf finding, not this bug. ## Asks 1. This is a **regression between 0.10.1 and 0.10.2** — please bisect that tag range (repro compiles on 0.10.1, fails on 0.10.2). Already fixed at head? 2. Is the ~1024 per-device limit intended, or the bug? Any flag to disable the `pipeline-start` async-wrapping that emits the invalid `async_computation`. ### System info (python version, jaxlib version, accelerator, etc.) **Environment:** jax 0.10.2, jaxlib 0.10.2, `jax[cuda12]` pip wheels, 8× H100, CUDA driver 580, NCCL 2.30.7. ### System info (`import jax; jax.print_environment_info()`, on the buggy 0.10.2) ``` jax: 0.10.2 jaxlib: 0.10.2 numpy: 2.5.1 python: 3.12.13 device info: NVIDIA H100-8, 8 local devices process_count: 1 platform: Linux x86_64 (kernel 6.13.2) $ nvidia-smi NVIDIA-SMI 580.82.07 Driver Version: 580.82.07 CUDA Version: 13.0 8x NVIDIA H100 (97 GB) ``` NCCL 2.30.7+cuda12.9 (bundled by the `jax[cuda12]` wheel). Installed via `pip install "jax[cuda12]==0.10.2"`.