Skip to content

Fusing Pallas Layout Transposes and Output Projections #449

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
prishajain1 wants to merge 1 commit into from
Closed

Fusing Pallas Layout Transposes and Output Projections #449

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 18 additions & 15 deletions src/maxdiffusion/kernels/custom_splash_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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, *_):
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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),
Expand All @@ -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
Expand Down
20 changes: 15 additions & 5 deletions src/maxdiffusion/models/attention_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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
Expand Down
1 change: 0 additions & 1 deletion src/maxdiffusion/tests/custom_splash_fixed_m_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Loading