diff --git a/src/maxdiffusion/kernels/custom_splash_attention.py b/src/maxdiffusion/kernels/custom_splash_attention.py index 1277200cf..f9bcf132c 100644 --- a/src/maxdiffusion/kernels/custom_splash_attention.py +++ b/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,7 +561,7 @@ 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), ] @@ -566,9 +569,9 @@ def v_index_map(h, i, j, *_): 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 diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 5a62e4e55..40d85be07 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/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 diff --git a/src/maxdiffusion/tests/custom_splash_fixed_m_test.py b/src/maxdiffusion/tests/custom_splash_fixed_m_test.py index 688d34e5f..c1a414172 100644 --- a/src/maxdiffusion/tests/custom_splash_fixed_m_test.py +++ b/src/maxdiffusion/tests/custom_splash_fixed_m_test.py @@ -98,7 +98,6 @@ def _run_kernel(self, q: jax.Array, k: jax.Array, v: jax.Array, use_fixed_m: boo use_fixed_m=use_fixed_m, ) out = kernel(q_in, k_in, v, mk) if use_fixed_m else kernel(q_in, k_in, v) - out = jnp.swapaxes(out, 1, 2) # (heads, dim, seq) -> (heads, seq, dim) return out.astype(jnp.float32), mk def test_online_matches_reference(self):