Trying wrap free ring schedule

Author: eltsai

src/maxdiffusion/generate_wan.py

@@ -299,6 +299,13 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
   max_logging.log(f"compile_time: {compile_time}")
   if writer and jax.process_index() == 0:
     writer.add_scalar("inference/compile_time", compile_time, global_step=0)
+  # Optional numerical dump of the decoded output for correctness checks
+  # (env WAN_DUMP_OUTPUT=/path/to.npy). Used to A/B attention variants.
+  if os.environ.get("WAN_DUMP_OUTPUT"):
+    import numpy as _np
+
+    _np.save(os.environ["WAN_DUMP_OUTPUT"], _np.asarray(videos[0], dtype=_np.float32))
+    max_logging.log(f"WAN_DUMP_OUTPUT: wrote {os.environ['WAN_DUMP_OUTPUT']}")
   saved_video_path = []
   for i in range(len(videos)):
     video_path = f"{filename_prefix}wan_output_{config.seed}_{i}.mp4"

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

src/maxdiffusion/models/attention_flax.py

@@ -766,24 +766,47 @@ def wrap_ulysses_attention(query, key, value):
   return x
 
 
-def _usp_groups(num_shards: int, ulysses_degree: int):
-  """Factorizes the `context` axis into a U x R (Ulysses x Ring) grid.
+INTERNAL_RING_AXIS = "ring"
+INTERNAL_ULYSSES_AXIS = "ulysses"
+
+
+def _replace_mesh_axis(axis_spec, old_axis, new_axes):
+  """Replace a mesh-axis name (or one nested inside a tuple) with new_axes."""
+  if axis_spec == old_axis:
+    return new_axes
+  if isinstance(axis_spec, tuple):
+    replacement = []
+    for axis in axis_spec:
+      if axis == old_axis:
+        replacement.extend(new_axes)
+      else:
+        replacement.append(axis)
+    return tuple(replacement)
+  return axis_spec
+
+
+def _replace_mesh_axis_names(axis_names, old_axis, new_axes):
+  return jax.sharding.PartitionSpec(*(_replace_mesh_axis(a, old_axis, new_axes) for a in axis_names))
 
-  Device index along the `context` axis is laid out u-major: `c = u*R + r`, so
-  - the Ulysses all-to-all group (fixed r, varying u) is `[u*R + r for u]`, and
-  - the Ring rotation group (fixed u, varying r) is `[u*R + r for r]`.
 
-  Returns `(U, R, ulysses_groups, ring_perm)` where `ulysses_groups` is the
-  `axis_index_groups` for the all-to-all and `ring_perm` is a `ppermute` perm
-  that rotates K/V by +1 *within each ring sub-group only*.
+def _create_internal_ulysses_ring_mesh(
+    mesh, ring_shards, ulysses_shards, ring_axis=INTERNAL_RING_AXIS, ulysses_axis=INTERNAL_ULYSSES_AXIS
+):
+  """Split the public `context` mesh axis into private (ring, ulysses) axes.
+
+  Ported from origin/main's tested implementation (commit c104db51). The reshape
+  is `(..., ring_shards, ulysses_shards, ...)` with the Ulysses axis INNERMOST, so
+  for ulysses_shards==2 the Ulysses group is consecutive device ids (the two cores
+  of one chip): the all-to-all stays intra-chip while the ring rotates across chips.
   """
-  U = ulysses_degree
-  if num_shards % U != 0:
-    raise ValueError(f"ulysses_degree={U} must divide context shard count {num_shards}.")
-  R = num_shards // U
-  ulysses_groups = [[u * R + r for u in range(U)] for r in range(R)]
-  ring_perm = [(u * R + r, u * R + (r + 1) % R) for u in range(U) for r in range(R)]
-  return U, R, ulysses_groups, ring_perm
+  mesh_axis_names = tuple(mesh.axis_names)
+  context_axis_index = mesh_axis_names.index("context")
+  devices = mesh.devices
+  new_shape = devices.shape[:context_axis_index] + (ring_shards, ulysses_shards) + devices.shape[context_axis_index + 1 :]
+  new_axis_names = (
+      mesh_axis_names[:context_axis_index] + (ring_axis, ulysses_axis) + mesh_axis_names[context_axis_index + 1 :]
+  )
+  return Mesh(devices.reshape(new_shape), new_axis_names)
 
 
 def _ulysses_ring_attention(
@@ -802,49 +825,66 @@ def _ulysses_ring_attention(
     use_base2_exp: bool = True,
     use_experimental_scheduler: bool = False,
 ) -> jax.Array:
-  """Hybrid Ulysses + Ring (USP) sequence parallelism with the custom kernel.
-
-  Splits the single `context` mesh axis into a U x R grid using collective
-  sub-groups (no extra mesh axes needed):
-    1. all-to-all WITHIN each Ulysses group (size U): trade sequence for heads,
-       so each device holds the full sequence of its ring-chunk but heads/U heads;
-    2. ring (ppermute) WITHIN each Ring group (size R): rotate K/V R times and
-       merge via online softmax, using the custom dense kernel per step;
+  """Hybrid Ulysses + Ring (USP) with the CUSTOM splash kernel on main's mesh.
+
+  Uses origin/main's explicit internal `(ring, ulysses)` mesh
+  (`_create_internal_ulysses_ring_mesh`, commit c104db51) instead of single-axis
+  collective sub-groups: the public `context` axis is reshaped with the Ulysses
+  axis innermost, so the Ulysses all-to-all stays INTRA-chip and the ring rotates
+  ACROSS chips. The per-shard attention is our custom splash kernel
+  (`make_custom_ring_attention`), not the tokamax_ring kernel main uses.
+
+    1. all-to-all over the (intra-chip) Ulysses axis: trade sequence for heads;
+    2. ring (full ppermute) over the (cross-chip) ring axis, online-softmax merge;
     3. all-to-all back to restore the sequence-sharded / full-heads layout.
 
-  U is read from the env var ULYSSES_RING_DEGREE (default 2); R = context // U.
-  U=context reduces to pure Ulysses, U=1 to pure Ring.
+  U = ULYSSES_RING_DEGREE (default 2); R = context // U. U=context -> pure
+  Ulysses, U=1 -> pure Ring (all on the same custom kernel).
   """
   axis_name = "context"
-  num_shards = mesh.shape[axis_name]
-  ulysses_degree = int(os.environ.get("ULYSSES_RING_DEGREE", "2"))
-  U, R, ulysses_groups, ring_perm = _usp_groups(num_shards, ulysses_degree)
+  num_context_shards = mesh.shape[axis_name]
+  num_ulysses_shards = int(os.environ.get("ULYSSES_RING_DEGREE", "2"))
+  if num_ulysses_shards <= 0 or num_context_shards % num_ulysses_shards != 0:
+    raise ValueError(
+        f"ULYSSES_RING_DEGREE={num_ulysses_shards} must be a positive divisor of "
+        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_shards)
-  key, _ = _reshape_data_for_flash(key, heads, num_shards)
-  value, _ = _reshape_data_for_flash(value, heads, num_shards)
+  query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_context_shards)
+  key, _ = _reshape_data_for_flash(key, heads, num_context_shards)
+  value, _ = _reshape_data_for_flash(value, heads, num_context_shards)
   num_heads = query.shape[1]
-  if num_heads % U != 0:
-    raise ValueError(f"Ulysses+Ring requires heads divisible by the Ulysses degree U={U}, got heads={num_heads}.")
+  if num_heads % num_ulysses_shards != 0:
+    raise ValueError(f"Ulysses+Ring requires heads divisible by U={num_ulysses_shards}, got heads={num_heads}.")
 
   bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes = _extract_custom_block_sizes(flash_block_sizes)
   if heads_per_tile > 1:
     raise NotImplementedError("ulysses_ring_custom currently supports heads_per_tile == 1 only.")
 
+  internal_mesh = _create_internal_ulysses_ring_mesh(mesh, num_ring_shards, num_ulysses_shards)
+  ring_axis = INTERNAL_RING_AXIS
+  ulysses_axis = INTERNAL_ULYSSES_AXIS
+
   q_axis_names = nn.logical_to_mesh_axes(axis_names_q)
   kv_axis_names = nn.logical_to_mesh_axes(axis_names_kv)
+  internal_q_axis_names = _replace_mesh_axis_names(q_axis_names, axis_name, (ring_axis, ulysses_axis))
+  internal_kv_axis_names = _replace_mesh_axis_names(kv_axis_names, axis_name, (ring_axis, ulysses_axis))
 
   @functools.partial(
       jax.shard_map,
-      mesh=mesh,
-      in_specs=(q_axis_names, kv_axis_names, kv_axis_names),
-      out_specs=q_axis_names,
+      mesh=internal_mesh,
+      in_specs=(internal_q_axis_names, internal_kv_axis_names, internal_kv_axis_names),
+      out_specs=internal_q_axis_names,
       check_vma=False,
   )
   def wrap_ulysses_ring_attention(query, key, value):
-    # (1) Ulysses all-to-all within each U-group: heads -> sequence swap, so each
-    # device now holds the full ring-chunk sequence with heads/U heads.
-    a2a = functools.partial(jax.lax.all_to_all, axis_name=axis_name, axis_index_groups=ulysses_groups, tiled=True)
+    # (1) Ulysses all-to-all over the (intra-chip) ulysses axis: heads -> sequence,
+    # so each device holds the full ring-chunk sequence with heads/U heads.
+    a2a = functools.partial(jax.lax.all_to_all, axis_name=ulysses_axis, tiled=True)
     query = a2a(query, split_axis=1, concat_axis=2)
     key = a2a(key, split_axis=1, concat_axis=2)
     value = a2a(value, split_axis=1, concat_axis=2)
@@ -857,7 +897,7 @@ def wrap_ulysses_ring_attention(query, key, value):
     value, _, _ = _pad_data_for_flash(value, heads, bkv)
 
     bsizes = custom_splash._BlockSizes(block_q=bq, block_kv=bkv, block_kv_compute=bkv_compute)
-    # (2) Ring over R within each U-group (rotation restricted by ring_perm).
+    # (2) Ring (full ppermute over the cross-chip ring axis) with the custom kernel.
     ring_kernel = tokamax_ring_attention_kernel.make_custom_ring_attention(
         block_sizes=bsizes,
         bkv_compute_in=bkv_compute_in,
@@ -866,19 +906,20 @@ def wrap_ulysses_ring_attention(query, key, value):
         use_base2_exp=use_base2_exp,
         use_experimental_scheduler=use_experimental_scheduler,
         vmem_limit_bytes=vmem_limit_bytes,
-        ring_axis=axis_name,
-        ring_size=R,
-        perm=ring_perm,
+        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)
     attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype)
 
-    # (3) Ulysses all-to-all back: sequence -> heads swap, restoring the layout.
+    # (3) Ulysses all-to-all back: sequence -> heads, restoring the layout.
     attention_output = a2a(attention_output, split_axis=2, concat_axis=1)
     return attention_output
 
   x = wrap_ulysses_ring_attention(query, key, value)
+  x = jax.lax.with_sharding_constraint(x, q_axis_names)
   x = x[:, :, :orig_q_seq_len, :]
   x = _reshape_heads_to_head_dim(x)
   return x