Dynamic-slice fusion hides Mosaic-GPU collective kernels from collective-memory coloring

[mcwitt]

When a Mosaic-GPU collective kernel consumes a slice of another op's output, DynamicSliceFusionRewriterV2 wraps the kernel and the slice into a fusion. The fusion boundary then prevents the collective-memory coloring from selecting the kernel's operands and results, so its buffers are allocated in default memory and NCCL symmetric window registration fails at the first execution.

Context: we hit this while porting a multi-process expert-parallel MoE transport (Pallas Mosaic-GPU kernels that write peer buffers through plgpu.remote_ref) to a full training step. It supersedes #47406: the failures we first attributed to allocator alignment are produced by this pass.

Minimal repro

Verified on GB200 (4 processes × 1 GPU, one node, stock nightly 0.11.1.dev20260809): with default flags, every process fails. With --xla_gpu_enable_dynamic_slice_fusion=false, every process prints ok. The first kernel runs unfused and works. The second kernel consumes (y1 @ w)[:15] and fails.

# failing arm (default flags), one process per GPU
for i in 0 1 2 3; do
  XLA_PYTHON_CLIENT_ALLOCATOR=cuda_async CUDA_VISIBLE_DEVICES=$i \
    python repro.py 127.0.0.1:9991 4 $i > /tmp/repro_fail_$i.log 2>&1 &
done; wait
# workaround arm
for i in 0 1 2 3; do
  XLA_PYTHON_CLIENT_ALLOCATOR=cuda_async CUDA_VISIBLE_DEVICES=$i \
  XLA_FLAGS=--xla_gpu_enable_dynamic_slice_fusion=false \
    python repro.py 127.0.0.1:9992 4 $i > /tmp/repro_ok_$i.log 2>&1 &
done; wait

repro.py

# python repro.py <coordinator_host:port> <num_processes> <process_id>  # 1 process/GPU
import sys

import jax
import jax.numpy as jnp
import numpy as np
from jax import lax
from jax.experimental import pallas as pl
from jax.experimental.pallas import mosaic_gpu as plgpu
from jax.sharding import AxisType, Mesh
from jax.sharding import PartitionSpec as P

coord, n_proc, pid = sys.argv[1], int(sys.argv[2]), int(sys.argv[3])
jax.distributed.initialize(coord, num_processes=n_proc, process_id=pid)
n = jax.device_count()
mesh = Mesh(np.array(jax.devices()), ("x",), axis_types=(AxisType.Explicit,))


def ring_shift(x):
    """Every device writes its block into the right neighbor's output."""

    def body(x_ref, out_ref):
        sem = pl.get_global(plgpu.SemaphoreType.REGULAR)
        dev = lax.axis_index("x")
        right = lax.rem(dev + 1, lax.axis_size("x"))
        dst = plgpu.remote_ref(out_ref, {"x": right})

        def scoped(smem, barrier):
            plgpu.copy_gmem_to_smem(x_ref, smem, barrier)
            plgpu.barrier_wait(barrier)
            plgpu.copy_smem_to_gmem(smem, dst)
            plgpu.wait_smem_to_gmem(0, wait_read_only=False)

        pl.run_scoped(scoped, plgpu.SMEM(x_ref.shape, x_ref.dtype), plgpu.Barrier(num_arrivals=1),
                      collective_axes="wg")
        pl.semaphore_signal(sem, device_id={"x": right})
        pl.semaphore_wait(sem, value=1)

    return plgpu.kernel(body, grid=(1,), grid_names=("sm",), num_threads=1, thread_name="wg",
                        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype))(x)


def f():
    dev = lax.axis_index("x").astype(jnp.float32)
    a = jnp.full((19, 256), dev, jnp.float32)
    y1 = ring_shift(a)
    # The second collective kernel consumes a static slice of a GEMM output.
    # The dynamic-slice fusion pass wraps the kernel and the slice into a
    # fusion, and the coloring no longer selects the kernel's buffers.
    w = jnp.full((256, 256), 0.01, jnp.float32)
    b = (y1 @ w)[:15]
    y2 = ring_shift(b)
    return y1, y2


run = jax.jit(jax.shard_map(f, mesh=mesh, in_specs=(), out_specs=(P("x"), P("x"))))
y1, y2 = jax.block_until_ready(run())
left = float((jax.process_index() - 1) % n)  # 1 process per GPU: global device rank == process index
left2 = float((jax.process_index() - 2) % n)
np.testing.assert_array_equal(np.asarray(y1.addressable_shards[0].data)[0, :4], np.full(4, left, np.float32))
np.testing.assert_allclose(np.asarray(y2.addressable_shards[0].data)[0, :4], np.full(4, left2 * 2.56, np.float32), rtol=1e-3)
print(f"process {jax.process_index()}: ok ({n} devices)", flush=True)

Actual

Every process fails at the first execution:

jax.errors.JaxRuntimeError: INTERNAL: NCCL operation nccl_status failed: unhandled cuda error ... 'Cuda failure 1 'invalid argument''. [executable_name='jit_f']

TF_CPP_VMODULE=nccl_symmetric_memory=3,collective_memory=3 identifies the failing registration. The collective arena registers fine, then registration is requested for the GEMM output allocation that passed through the fusion boundary — a default-memory (cudaMallocAsync pool) pointer:

nccl_symmetric_memory.cc:85] Create NCCL symmetric memory on comm=... from: ptr=0xf199f4a00000; size=2097152
collective_memory.cc:560] [0]  Acquire collective memory for global device id 0: run_id=... symmetric=1 multicast=0 peer=0
nccl_symmetric_memory.cc:85] Create NCCL symmetric memory on comm=... from: ptr=0x340000000; size=19456
→ Cuda failure 1 'invalid argument'

(19456 bytes = the full f32[19,256] GEMM output; the fusion operand is the whole producer buffer, sliced inside the fusion. With a BFC default arena instead of cudaMallocAsync, the same escape surfaces as Window address must be suitably aligned.)

The after-optimizations HLO shows the mechanism. With default flags, the second kernel sits inside the fusion and no operand or result carries S(1):

%dynamic-slice-fusion (p0: f32[19,256], p1: s32[1]) -> (f32[15,256], s32[1]) {
  ...
  ROOT %mpmd_map.16 = (f32[15,256]{1,0}, s32[1]{0}) custom-call(%slice.4, %p1), custom_call_target="mosaic_gpu_v2", ...
}
...
%dynamic_slice_fusion = (f32[15,256]{1,0}, s32[1]{0}) fusion(%gemm_fusion_dot_general.3, %copy.8), kind=kCustom, calls=%dynamic-slice-fusion, ...

while the unfused first kernel in the same module is fully colored (%loop_broadcast_fusion = f32[19,256]{1,0:S(1)} feeding %mpmd_map.14). With the flag off, both call results are colored:

%mpmd_map.14 = (f32[19,256]{1,0:S(1)}, s32[1]{0:S(1)}) custom-call(...), custom_call_target="mosaic_gpu_v2", ...
%mpmd_map.15 = (f32[15,256]{1,0:S(1)}, s32[1]{0:S(1)}) custom-call(%wrapped_slice, ...), custom_call_target="mosaic_gpu_v2", ...

Expected

Fusing a collective custom call does not change the memory space of its operands and results, or the rewriter does not wrap such calls.

Root cause hypothesis

The rewriter's predicate accepts any registered typed-FFI custom call, with no exclusion for calls that need collective memory. The collective coloring and copy-insertion analyses (for example IsMosaicWithMultimem and its callers) do not look through the fusion boundary, and the fusion instruction carries neither the custom-call opcode nor the hero's backend config. At execution, JAX's Mosaic-GPU FFI prepare handler (MosaicGpuPrepare) requests symmetric registration of the resulting default-memory buffer, which fails as shown above. The pass is default-on since PR #43831 (2026-06-30).

Workaround

--xla_gpu_enable_dynamic_slice_fusion=false. Validated at 4, 8, and 64 ranks; no measurable performance cost on our workload.

Environment

• GB200 NVL72, one node, one process per GPU. CUDA 13, driver 595.71.05, NCCL 2.30.7.
• jax, jaxlib, jax-cuda13-plugin, jax-cuda13-pjrt 0.11.1.dev20260809 (stock nightlies from the jax-public-nightly registry); also reproduced on .dev20260816.
• Analyzed at xla 60f8069e8b. The code paths are unchanged at f0bfd165e0 (2026-08-17).
• Related: RaggedAllToAll one-shot path crashes with CUDA_ERROR_ILLEGAL_ADDRESS above 32 ranks #47283. Supersedes NCCL symmetric window registration fails 4096-byte alignment when collective memory uses a separate BFC allocator #47406.

[akuegel]

@ezhulenev maybe you have an idea how we can possibly filter out such custom calls that require collective memory?

[ezhulenev]

Hm, I think instead of filtering it's better for fix coloring and propagate it through the fusion boundary. I see there is some ongoing work with putting collectives into fusions (targeting Triton?). it will hit the same problem, if collective codegen will require symmetric memory.