Optimize WAN attention by fusing Pallas transposes and Output Projection contractions

Author: prishajain1

src/maxdiffusion/kernels/custom_splash_attention.py

@@ -237,12 +237,12 @@ def end():
     l = l_scratch_ref[...]
     if fuse_reciprocal:
       l_inv = jnp.tile(1.0 / l, (head_dim_v_repeats, 1))
-      o_ref[...] = (o_scratch_ref[...] * l_inv).astype(o_ref.dtype)
+      o_ref[...] = (o_scratch_ref[...] * l_inv).T.astype(o_ref.dtype)
     else:
       # Ring path: emit the un-normalized numerator plus the running softmax
       # stats (max logit `m` and linear denominator `l`) so the outer ring loop
       # can merge shard contributions and normalize only once at the very end.
-      o_ref[...] = o_scratch_ref[...].astype(o_ref.dtype)
+      o_ref[...] = o_scratch_ref[...].T.astype(o_ref.dtype)
     if l_ring_ref is not None:
       l_ring_ref[...] = l.astype(l_ring_ref.dtype)
     if m_ring_ref is not None:
@@ -399,7 +399,7 @@ def end():
     for h_local in range(heads_per_tile):
       l = l_scratch_ref[h_local]
       l_inv = jnp.tile(1.0 / l, (head_dim_v_repeats, 1))
-      o_ref[h_local] = (o_scratch_ref[h_local] * l_inv).astype(o_ref.dtype)
+      o_ref[h_local] = (o_scratch_ref[h_local] * l_inv).T.astype(o_ref.dtype)
 
 
 def _splash_attention_forward(
@@ -436,7 +436,7 @@ def q_index_map(h, i, j, *_):
     return (h, i, 0)
 
   def out_index_map(h, i, j, *_):
-    return h, 0, i
+    return h, i, 0
 
   def k_index_map(h, i, j, *_):
     return (h // q_heads_per_kv_head, j, 0)
@@ -453,13 +453,13 @@ def v_index_map(h, i, j, *_):
       jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32),
       jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32),
       jax.ShapeDtypeStruct((head_dim_v, bq), jnp.float32),
-      jax.ShapeDtypeStruct((num_q_heads, head_dim_v, actual_q_seq_len), q.dtype),
+      jax.ShapeDtypeStruct((num_q_heads, actual_q_seq_len, head_dim_v), q.dtype),
   ]
   out_specs = [
       pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)),
       pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)),
       pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)),
-      pl.BlockSpec((None, head_dim_v, bq), out_index_map),
+      pl.BlockSpec((None, bq, head_dim_v), out_index_map),
   ]
   grid_width = (actual_kv_seq_len + bkv - 1) // bkv
   grid_height = (actual_q_seq_len + bq - 1) // bq
@@ -540,7 +540,10 @@ def _splash_attention_forward_ring(
   def q_index_map(h, i, j, *_):
     return (h, i, 0)
 
-  def out_index_map(h, i, j, *_):
+  def out_index_map_o(h, i, j, *_):
+    return h, i, 0
+
+  def out_index_map_stats(h, i, j, *_):
     return h, 0, i
 
   def k_index_map(h, i, j, *_):
@@ -558,17 +561,17 @@ def v_index_map(h, i, j, *_):
       jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32),
       jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32),
       jax.ShapeDtypeStruct((head_dim_v, bq), jnp.float32),
-      jax.ShapeDtypeStruct((num_q_heads, head_dim_v, actual_q_seq_len), jnp.float32),
+      jax.ShapeDtypeStruct((num_q_heads, actual_q_seq_len, head_dim_v), jnp.float32),
       jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, actual_q_seq_len), jnp.float32),
       jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, actual_q_seq_len), jnp.float32),
   ]
   out_specs = [
       pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)),
       pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)),
       pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)),
-      pl.BlockSpec((None, head_dim_v, bq), out_index_map),
-      pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map),
-      pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map),
+      pl.BlockSpec((None, bq, head_dim_v), out_index_map_o),
+      pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map_stats),
+      pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map_stats),
   ]
   grid_width = (actual_kv_seq_len + bkv - 1) // bkv
   grid_height = (actual_q_seq_len + bq - 1) // bq
@@ -612,7 +615,7 @@ def v_index_map(h, i, j, *_):
       ),
       out_shape=out_shapes,
   )(mk, q, k, v)
-  out = jnp.swapaxes(all_out[3], 1, 2)  # (h, head_dim_v, s) -> (h, s, head_dim_v)
+  out = all_out[3]
   l = all_out[4][:, 0, :]  # (h, s)
   m = all_out[5][:, 0, :]  # (h, s)
   return out, m, l
@@ -653,7 +656,7 @@ def v_index_map(h, i, j, *_):
     return (h, j, 0)
 
   def out_index_map(h, i, j, *_):
-    return (h, 0, i)
+    return (h, i, 0)
 
   in_specs = [
       pl.BlockSpec((hpt, bq, head_dim_qk), q_index_map),
@@ -664,13 +667,13 @@ def out_index_map(h, i, j, *_):
       jax.ShapeDtypeStruct((hpt, NUM_SUBLANES, bq), jnp.float32),
       jax.ShapeDtypeStruct((hpt, NUM_SUBLANES, bq), jnp.float32),
       jax.ShapeDtypeStruct((hpt, head_dim_v, bq), jnp.float32),
-      jax.ShapeDtypeStruct((num_q_heads, head_dim_v, actual_q_seq_len), q.dtype),
+      jax.ShapeDtypeStruct((num_q_heads, actual_q_seq_len, head_dim_v), q.dtype),
   ]
   out_specs = [
       pl.BlockSpec((hpt, NUM_SUBLANES, bq), lambda *_: (0, 0, 0)),
       pl.BlockSpec((hpt, NUM_SUBLANES, bq), lambda *_: (0, 0, 0)),
       pl.BlockSpec((hpt, head_dim_v, bq), lambda *_: (0, 0, 0)),
-      pl.BlockSpec((hpt, head_dim_v, bq), out_index_map),
+      pl.BlockSpec((hpt, bq, head_dim_v), out_index_map),
   ]
   grid_width = (actual_kv_seq_len + bkv - 1) // bkv
   grid_height = (actual_q_seq_len + bq - 1) // bq

src/maxdiffusion/models/attention_flax.py

@@ -826,7 +826,6 @@ def wrap_ulysses_attention(query, key, value):
       else:
         vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0))
         attention_output = vmapped_splash(query, key, value)
-      attention_output = jnp.swapaxes(attention_output, 2, 3)
       attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype)
     else:
       # Run the same local splash kernel as standard TPU flash attention, but now
@@ -907,6 +906,8 @@ def wrap_ulysses_attention(query, key, value):
   if fold_batch:
     x = x.reshape(batch, num_heads, *x.shape[2:])
   x = x[:, :, :orig_q_seq_len, :]
+  if use_custom_kernel:
+    return x
   x = _reshape_heads_to_head_dim(x)
 
   return x
@@ -1218,9 +1219,9 @@ def wrap_ulysses_ring_attention(query, key, value):
           use_fixed_m=use_fixed_m,
       )
       if use_fixed_m:
-        attention_output = jnp.swapaxes(jax.vmap(splash_kernel, in_axes=(0, 0, 0, None))(query, key, value, mk_arr), 2, 3)
+        attention_output = jax.vmap(splash_kernel, in_axes=(0, 0, 0, None))(query, key, value, mk_arr)
       else:
-        attention_output = jnp.swapaxes(jax.vmap(splash_kernel, in_axes=(0, 0, 0))(query, key, value), 2, 3)
+        attention_output = jax.vmap(splash_kernel, in_axes=(0, 0, 0))(query, key, value)
     else:
       # (2b) Ring (full ppermute over the cross-chip ring axis) with the custom kernel.
       # bidirectional=True -> wrap-free schedule (streams K/V both directions one hop
@@ -1257,7 +1258,6 @@ def 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
 
 
@@ -2460,7 +2460,17 @@ def __call__(
     attn_output = checkpoint_name(attn_output, "attn_output")
 
     with jax.named_scope("proj_attn"):
-      hidden_states = self.proj_attn(attn_output)
+      if attn_output.ndim == 4:
+        kernel = self.proj_attn.kernel.value.reshape(self.heads, self.dim_head, -1)
+        hidden_states = jax.lax.dot_general(
+            attn_output,
+            kernel,
+            (((1, 3), (0, 1)), ((), ())),
+        )
+        if self.proj_attn.bias is not None:
+          hidden_states = hidden_states + self.proj_attn.bias.value
+      else:
+        hidden_states = self.proj_attn(attn_output)
       if self.drop_out.rate > 0:
         hidden_states = self.drop_out(hidden_states, deterministic=deterministic, rngs=rngs)
     return hidden_states