bring back wrap free ring for more tests

Author: eltsai

src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py

@@ -726,6 +726,109 @@ def make_ring_attention(
 # ---------------------------------------------------------------------------
 
 
+def _custom_bidirectional_ring_forward(
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    *,
+    block_sizes: "custom_splash._BlockSizes",
+    bkv_compute_in: int,
+    orig_q_seq_len: int,
+    orig_kv_seq_len: int,
+    use_base2_exp: bool,
+    use_experimental_scheduler: bool,
+    vmem_limit_bytes: int | None,
+    mask_value: float,
+    ring_axis: str,
+) -> jax.Array:
+  """Wrap-free (bidirectional) ring attention for a NON-wrapping ring axis.
+
+  On a torus dimension the +1-mod-R ppermute is nearest-neighbor, but on a
+  cut/non-wrapping axis (e.g. the size-4 z line of a v7x-16 slice) the wrap edge
+  (R-1 -> 0) spans the whole line diameter. Instead of one rotating stream with
+  that long, congested wrap, stream K/V BOTH directions one hop at a time:
+    - rightward stream: device i holds KV_{i-t} after t hops,
+    - leftward  stream: device i holds KV_{i+t} after t hops,
+  with out-of-range shards (line ends) masked out of the online softmax. Every
+  step is a single hop and uses both link directions; no edge ever spans the
+  diameter.
+
+  Trade-off: each device computes ~2x attention blocks (the line-end ones are
+  masked), traded for the removed multi-hop wrap. Net win when the ring is
+  comms-bound (the case on a non-wrapping axis). Operates on the FULL real ring
+  axis (no sub-group perm).
+  """
+  axis_size = lax.axis_size(ring_axis)
+  idx = lax.axis_index(ring_axis)
+  exp_fn = jnp.exp2 if use_base2_exp else jnp.exp
+
+  def _attn(kc, vc):
+    o, m, l = custom_splash._splash_attention_forward_ring(  # pylint: disable=protected-access
+        q,
+        kc,
+        vc,
+        block_sizes,
+        bkv_compute_in,
+        q_seq_len=orig_q_seq_len,
+        kv_seq_len=orig_kv_seq_len,
+        use_base2_exp=use_base2_exp,
+        use_experimental_scheduler=use_experimental_scheduler,
+        vmem_limit_bytes=vmem_limit_bytes,
+    )
+    return o.astype(jnp.float32), m.astype(jnp.float32), l.astype(jnp.float32)
+
+  def _merge(m, l, o, mc, lc, oc, valid):
+    # Nullify invalid (line-end) contributions. Force mc to mask_value (beta -> 0)
+    # AND zero lc/oc so a non-finite zero-buffer result can't leak via 0*inf=nan.
+    mc = jnp.where(valid, mc, mask_value)
+    lc = jnp.where(valid, lc, 0.0)
+    oc = jnp.where(valid, oc, 0.0)
+    m_next = jnp.maximum(m, mc)
+    alpha = exp_fn(m - m_next)
+    beta = exp_fn(mc - m_next)
+    return m_next, alpha * l + beta * lc, alpha[..., None] * o + beta[..., None] * oc
+
+  # t=0: own shard (always valid). _attn returns (o, m, l).
+  o, m, l = _attn(k, v)
+
+  # Non-wrapping one-hop shifts (line ends send/receive nothing).
+  shift_r = partial(lax.ppermute, axis_name=ring_axis, perm=[(i, i + 1) for i in range(axis_size - 1)])
+  shift_l = partial(lax.ppermute, axis_name=ring_axis, perm=[(i, i - 1) for i in range(1, axis_size)])
+
+  # Prime buffers for t=1 (one hop each direction): device i -> KV_{i-1}, KV_{i+1}.
+  kr, vr = shift_r(k), shift_r(v)
+  kl, vl = shift_l(k), shift_l(v)
+
+  def body(carry, t):
+    m, l, o, kr, vr, kl, vl = carry
+    valid_r = (idx - t) >= 0
+    valid_l = (idx + t) <= (axis_size - 1)
+    # Feed real (own) K/V on invalid steps so _attn never runs on a degenerate
+    # zero buffer (line ends receive 0 from the partial ppermute); masked below.
+    kr_s, vr_s = jnp.where(valid_r, kr, k), jnp.where(valid_r, vr, v)
+    kl_s, vl_s = jnp.where(valid_l, kl, k), jnp.where(valid_l, vl, v)
+    # Compute against the current shards (KV_{i-t}, KV_{i+t}) ...
+    o_r, m_r, l_r = _attn(kr_s, vr_s)
+    m, l, o = _merge(m, l, o, m_r, l_r, o_r, valid_r)
+    o_l, m_l, l_l = _attn(kl_s, vl_s)
+    m, l, o = _merge(m, l, o, m_l, l_l, o_l, valid_l)
+    # ... and prefetch the next hop (independent of the matmuls above -> overlaps).
+    kr_n, vr_n = shift_r(kr), shift_r(vr)
+    kl_n, vl_n = shift_l(kl), shift_l(vl)
+    return (m, l, o, kr_n, vr_n, kl_n, vl_n), None
+
+  (_, l_final, o_final, *_), _ = lax.scan(
+      body,
+      (m, l, o, kr, vr, kl, vl),
+      xs=jnp.arange(1, axis_size),
+      length=axis_size - 1,
+      unroll=True,
+  )
+
+  l_inv = jnp.where(l_final == 0.0, 0.0, 1.0 / l_final)
+  return (o_final * l_inv[..., None]).astype(q.dtype)
+
+
 def _custom_ring_attention_forward(
     q: jax.Array,
     k: jax.Array,
@@ -742,6 +845,7 @@ def _custom_ring_attention_forward(
     ring_axis: str,
     ring_size: int | None = None,
     perm: list[tuple[int, int]] | None = None,
+    bidirectional: bool = False,
 ) -> jax.Array:
   """Forward-only ring attention using the custom dense splash kernel.
 
@@ -776,6 +880,26 @@ def _custom_ring_attention_forward(
     Normalized attention output, shape `(num_q_heads, q_seq_len, head_dim_v)`.
   """
   axis_size = lax.axis_size(ring_axis)
+  if bidirectional:
+    if perm is not None or (ring_size is not None and ring_size != axis_size):
+      raise ValueError(
+          "bidirectional (wrap-free) ring requires perm=None and ring_size==axis_size "
+          "(it operates on the full real ring axis)."
+      )
+    return _custom_bidirectional_ring_forward(
+        q,
+        k,
+        v,
+        block_sizes=block_sizes,
+        bkv_compute_in=bkv_compute_in,
+        orig_q_seq_len=orig_q_seq_len,
+        orig_kv_seq_len=orig_kv_seq_len,
+        use_base2_exp=use_base2_exp,
+        use_experimental_scheduler=use_experimental_scheduler,
+        vmem_limit_bytes=vmem_limit_bytes,
+        mask_value=mask_value,
+        ring_axis=ring_axis,
+    )
   if ring_size is None:
     ring_size = axis_size
   if perm is None:
@@ -847,6 +971,7 @@ def make_custom_ring_attention(
     ring_axis: str = "context",
     ring_size: int | None = None,
     perm: list[tuple[int, int]] | None = None,
+    bidirectional: bool = False,
 ):
   """Builds a forward-only ring-attention callable around the custom kernel.
 
@@ -858,6 +983,10 @@ def make_custom_ring_attention(
   `ring_size` / `perm` let a caller restrict the rotation to a ring sub-group of
   the axis (for the hybrid Ulysses+Ring / USP split); when omitted the rotation
   covers the whole `ring_axis`.
+
+  `bidirectional=True` selects the wrap-free schedule (streams K/V both directions
+  one hop at a time) for a NON-wrapping ring axis, avoiding the diameter-length
+  wrap hop. Requires `perm=None` and the full real ring axis (no sub-group).
   """
 
   def _ring(q, k, v):
@@ -876,6 +1005,7 @@ def _ring(q, k, v):
         ring_axis=ring_axis,
         ring_size=ring_size,
         perm=perm,
+        bidirectional=bidirectional,
     )
 
   return _ring