Default v5e-8 mesh not efficient?

[sssshhhhhh]

On v5e-8 2x4 topology I get these device coords which are sequential already so the mesh ids are directly

jax/jax/_src/mesh_utils.py

Line 72 in b08a105

_V5E_TRAY_RING_ORDER = (0, 1, 2, 3, 7, 6, 5, 4)

TPU_0(process=0,(0,0,0,0))
TPU_1(process=0,(1,0,0,0))
TPU_2(process=0,(0,1,0,0))
TPU_3(process=0,(1,1,0,0))
TPU_4(process=0,(0,2,0,0))
TPU_5(process=0,(1,2,0,0))
TPU_6(process=0,(0,3,0,0))
TPU_7(process=0,(1,3,0,0))
jax.make_mesh((8,), ('x',)).device_ids
>>> [0 1 2 3 7 6 5 4]

Which made me ask why 1(1,0)/2(0, 1); or 3(1,1)/7(1,3) are adjacent. So I implemented a bidirectional all gather
(this img basically https://jax-ml.github.io/scaling-book/assets/img/all-gather.gif)

import functools
import jax
from jax import lax
from jax import numpy as jnp
from jax.experimental import pallas as pl
from jax.experimental.pallas import tpu as pltpu
from jax.sharding import PartitionSpec as P, NamedSharding, Mesh
import numpy as np

num_devices = jax.local_device_count()
partition = P('x', None)
mesh = jax.make_mesh((num_devices,), ('x',))
devices = jax.local_devices()
# mesh = Mesh(np.asarray(devices)[[0,2,4,6,7,5,3,1]], ('x',))
sharding = NamedSharding(mesh, partition)
input_arr = jax.random.uniform(jax.random.key(0), (8 * num_devices, 2**22), jnp.float32)
input_arr = jax.device_put(input_arr, sharding)

def mod(x, n):
    return lax.rem(x + n, n)

def all_gather_kernel(input_ref, output_ref, local_copy_sem, send_sem, recv_sems):
    step = pl.program_id(0)
    my_id = lax.axis_index('x')
    left_neighbor = mod(my_id - 1, num_devices)
    right_neighbor = mod(my_id + 1, num_devices)
    left_copy_slot = mod(my_id + step, num_devices)
    right_copy_slot = mod(my_id - step, num_devices)

    @pl.when(step == 0)
    def _():
        local_copy = pltpu.async_copy(input_ref, output_ref.at[my_id], local_copy_sem)
        barrier_sem = pltpu.get_barrier_semaphore()
        pltpu.semaphore_signal(barrier_sem, device_id=(left_neighbor,))
        pltpu.semaphore_signal(barrier_sem, device_id=(right_neighbor,))
        pltpu.semaphore_wait(barrier_sem, 2)
        local_copy.wait()

    left_copy = pltpu.make_async_remote_copy(
        src_ref=output_ref.at[left_copy_slot],
        dst_ref=output_ref.at[left_copy_slot],
        send_sem=send_sem.at[0],
        recv_sem=recv_sems.at[0, step],
        device_id=(left_neighbor,),
        device_id_type=pltpu.DeviceIdType.MESH,
    )
    right_copy = pltpu.make_async_remote_copy(
        src_ref=output_ref.at[right_copy_slot],
        dst_ref=output_ref.at[right_copy_slot],
        send_sem=send_sem.at[1],
        recv_sem=recv_sems.at[1, step],
        device_id=(right_neighbor,),
        device_id_type=pltpu.DeviceIdType.MESH,
    )
    left_copy.start()

    @pl.when(step < pl.num_programs(0) - 1)
    def _():
        right_copy.start()
        right_copy.wait()
    left_copy.wait()


@functools.partial(jax.jit, in_shardings=sharding, out_shardings=NamedSharding(mesh, P()))
@functools.partial(jax.shard_map, mesh=mesh, in_specs=partition, out_specs=P())
def gather(input_arr):
    out_shape = jax.ShapeDtypeStruct((num_devices, 8, input_arr.shape[-1]), jnp.float32, vma=frozenset())
    grid_spec = pltpu.PrefetchScalarGridSpec(
        num_scalar_prefetch=0,
        in_specs=[pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY)],
        out_specs=pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),
        scratch_shapes=([
            pltpu.SemaphoreType.DMA,
            pltpu.SemaphoreType.DMA((2,)),
            pltpu.SemaphoreType.DMA((2, num_devices//2,)),
        ]),
        grid=(num_devices//2,),
    )
    return pl.pallas_call(
        all_gather_kernel,
        out_shape=out_shape,
        grid_spec=grid_spec,
    )

gather(input_arr).block_until_ready()

%timeit gather(input_arr).block_until_ready()
>>> 24.9 ms ± 49 μs per loop(mean ± std. dev. of 7 runs, 10 loops each)

This is double the roofline but if I change the mesh to this it's a lot faster

4 steps each sends 8*2^22*4=134,217,728 B/step
ICI link: 4.5e10 B/s
134,217,728/4.5e10*4*1000=11.93 ms
mesh = Mesh(np.asarray(devices)[[0,2,4,6,7,5,3,1]], ('x',))
>>> 16.4 ms ± 53.8 μs per loop(mean ± std. dev. of 7 runs, 100 loops each)

Am I calculating the neighbors wrong? Is the mesh supposed to be like this?