test(splash-ring): use full mesh and all-gather in mhpt invariance test

Author: edgexyz

src/maxdiffusion/kernels/splash_attention/ring_attention_kernel_test.py

@@ -21,6 +21,7 @@
 from absl.testing import parameterized
 import jax
 from jax import random
+from jax.experimental import multihost_utils
 import jax.numpy as jnp
 import numpy as np
 from . import base
@@ -180,13 +181,17 @@ def setUp(self):
       dtype=[jnp.bfloat16],
   )
   def test_heads_per_tile_matches_single_head(self, heads_per_tile, head_dim, dtype):
-    ring_size = 2
-    num_heads = 4  # MHA (num_q_heads == num_kv_heads); divisible by heads_per_tile.
-    if len(jax.devices()) < ring_size:
-      self.skipTest(f"This test requires {ring_size} devices, but has only {len(jax.devices())}.")
+    # Use ALL devices so the sharded outputs are addressable on every process
+    # (a subset mesh like jax.devices()[:2] lives only on process 0, making the
+    # result non-addressable on the other hosts). jax.devices() is itself a
+    # global barrier, so every host must launch this test together regardless.
+    ring_size = jax.device_count()
+    num_heads = 8  # MHA (num_q_heads == num_kv_heads); divisible by heads_per_tile.
+    if ring_size < 2:
+      self.skipTest(f"This test needs at least 2 devices, but has {ring_size}.")
 
     ring_axis = "ring"
-    devices = np.asarray(jax.devices()[:ring_size]).reshape(1, ring_size)
+    devices = np.asarray(jax.devices()).reshape(1, ring_size)
     mesh = jax.sharding.Mesh(devices, ("heads", ring_axis))
     seq_len = 1024 * ring_size
 
@@ -236,10 +241,11 @@ def ring_attn(ring_kernel, q, k, v, segment_ids):
     out_mhpt = run(heads_per_tile)  # multi-head-per-tile (flash_attention_kernel_mhpt)
 
     # Pure tiling => numerically equivalent to the single-head-per-tile baseline.
-    # The ring mesh is jax.devices()[:ring_size] (all on process 0), so the
-    # outputs are only addressable there; use the multi-controller-safe compare
-    # so this passes on every host, not just the owner.
-    self.assert_allclose_mcjax(out_mhpt, out_ref, rtol=5e-3, atol=5e-3)
+    # Outputs are sharded across all hosts; all-gather to a fully-replicated host
+    # array on every process, then compare with the standard helper.
+    out_mhpt = multihost_utils.process_allgather(out_mhpt, tiled=True)
+    out_ref = multihost_utils.process_allgather(out_ref, tiled=True)
+    self._assert_allclose(out_mhpt, out_ref, rtol=5e-3, atol=5e-3)
 
 
 if __name__ == "__main__":

src/maxdiffusion/kernels/splash_attention/splash_attention_test_utils.py

@@ -16,7 +16,6 @@
 import unittest
 from absl.testing import parameterized
 import jax
-from jax.experimental import multihost_utils
 import jax.numpy as jnp
 import numpy as np
 
@@ -77,43 +76,6 @@ def _assert_allclose(self, x, y, **kwargs):
     self.assertTupleEqual(x.shape, y.shape)
     np.testing.assert_allclose(x, y, **kwargs)
 
-  def assert_allclose_mcjax(self, x, y, *, rtol, atol):
-    """`allclose` that is safe under multi-controller (multi-host) JAX.
-
-    Some tests build their device mesh from a subset of the global devices --
-    e.g. `jax.devices()[:ring_size]`, which all live on process 0 -- so the
-    result `jax.Array`s are only addressable on that one process. Pulling them
-    to host with `np.testing.assert_allclose` (as `_assert_allclose` does)
-    raises `RuntimeError: ... spans non-addressable devices` on every *other*
-    process, failing the test on all but one host.
-
-    Instead, evaluate the comparison on-device (works on all processes, no host
-    fetch), read the scalar verdict only on the owning process (process 0, which
-    holds `jax.devices()[:ring_size]`), and broadcast it to every process with a
-    single collective. Every process runs the same two `broadcast_one_to_all`
-    calls in the same order, so there is no collective-participation mismatch /
-    deadlock. The result is identical on all hosts and genuinely reflects the
-    owner's computation.
-
-    Only use this for tests whose mesh is a subset of one process's devices; for
-    fully-sharded (every-process-owns-a-shard) arrays use `_assert_allclose`.
-    """
-    if x.dtype == np.dtype(jnp.bfloat16):
-      x = x.astype(jnp.float32)
-    if y.dtype == np.dtype(jnp.bfloat16):
-      y = y.astype(jnp.float32)
-    self.assertTupleEqual(x.shape, y.shape)
-    ok = jnp.all(jnp.abs(x - y) <= atol + rtol * jnp.abs(y))
-    max_err = jnp.max(jnp.abs(x - y))
-    is_owner = jax.process_index() == 0
-    local_ok = np.asarray(ok) if is_owner else np.array(True)
-    local_err = np.asarray(max_err) if is_owner else np.array(0.0, np.float32)
-    global_err = float(multihost_utils.broadcast_one_to_all(local_err))
-    self.assertTrue(
-        bool(multihost_utils.broadcast_one_to_all(local_ok)),
-        f"arrays differ: max abs err {global_err:.3e} exceeds rtol={rtol} atol={atol}",
-    )
-
 
 def create_segment_ids(seq_len: int, num_breaks: int = 2) -> base.SegmentIds:
   break_indices = np.random.choice(range(1, seq_len), num_breaks, replace=False)