revert wrap free ring, negative results

Author: eltsai

src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py

@@ -726,109 +726,6 @@ 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,
@@ -845,7 +742,6 @@ 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.
 
@@ -880,26 +776,6 @@ 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:
@@ -971,7 +847,6 @@ 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.
 
@@ -983,10 +858,6 @@ 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):
@@ -1005,7 +876,6 @@ def _ring(q, k, v):
         ring_axis=ring_axis,
         ring_size=ring_size,
         perm=perm,
-        bidirectional=bidirectional,
     )
 
   return _ring

src/maxdiffusion/models/attention_flax.py

@@ -850,9 +850,6 @@ def _ulysses_ring_attention(
         f"context shard count {num_context_shards}."
     )
   num_ring_shards = num_context_shards // num_ulysses_shards
-  # Wrap-free (bidirectional) ring schedule for a non-wrapping ring axis (e.g. the
-  # cut size-4 z line of a v7x-16 slice). Set USP_WRAP_FREE_RING=1 to A/B it.
-  wrap_free_ring = os.environ.get("USP_WRAP_FREE_RING", "0") == "1"
 
   query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_context_shards)
   key, _ = _reshape_data_for_flash(key, heads, num_context_shards)
@@ -908,7 +905,6 @@ def wrap_ulysses_ring_attention(query, key, value):
         vmem_limit_bytes=vmem_limit_bytes,
         ring_axis=ring_axis,
         ring_size=num_ring_shards,
-        bidirectional=wrap_free_ring,
     )
     vmapped_ring = jax.vmap(ring_kernel, in_axes=(0, 0, 0))
     attention_output = vmapped_ring(query, key, value)