Skip to content

Mhc optimization: Rewrites einsums for better xla fusion #4664

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

Open
h1t35h wants to merge 1 commit into
base: main
Choose a base branch
from
+12 −16
Open

Mhc optimization: Rewrites einsums for better xla fusion #4664

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
28 changes: 12 additions & 16 deletions src/maxtext/layers/mhc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,13 +67,10 @@ def sinkhorn(t, iters=20):
t = jax.nn.softmax(t, axis=-1) + eps
t = t / (jnp.sum(t, axis=-2, keepdims=True) + eps)

def body_fun(i, val):
val = val / (jnp.sum(val, axis=-1, keepdims=True) + eps)
val = val / (jnp.sum(val, axis=-2, keepdims=True) + eps)
return val
for _ in range(iters - 1):
t = t / (jnp.sum(t, axis=-1, keepdims=True) + eps)
t = t / (jnp.sum(t, axis=-2, keepdims=True) + eps)

# Use lax.fori_loop for an efficient, JIT-friendly loop
t = jax.lax.fori_loop(0, iters - 1, body_fun, t)
return t.astype(initial_dtype)


Expand Down Expand Up @@ -257,8 +254,9 @@ def __call__(
# x shape: [batch, seq, expansion_rate, emb]
b, s, k, d = x.shape

# 1. Flatten the tensor, and RMS normalization
norm_x = self.mhc_norm(jnp.reshape(x, (b, s, k * d)))
with jax.named_scope("mhc_norm"):
# 1. Flatten the tensor, and RMS normalization
norm_x = self.mhc_norm(jnp.reshape(x, (b, s, k * d)))

# 2. Pre mapping
pre_mapping = self.mapping(
Expand All @@ -269,7 +267,8 @@ def __call__(
1.0,
eps=1e-6,
)
layer_input = jnp.einsum("bskd,bsk -> bsd", x, pre_mapping, precision=self.matmul_precision)
# bskd, bsk -> bsd
layer_input = jnp.sum(x * jnp.expand_dims(pre_mapping, axis=3), axis=2)

@NuojCheng NuojCheng Jul 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

for my understanding, do you see perf gain by this change? I thought einsum would be more performant.

@h1t35h h1t35h Jul 30, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, even I found it a bit counter-intuitive einsum being more performant is what my understanding was as well. I think while being generally true for some cases like this one it seems better optimization is coming by moving away from it. Specifically, I can see the all-reduce overlap. I tested the changes individually by just changing the einsums and for_iloop but that does not seem to work I'm only seeing the All-Reduce overlap when both the changes are present.

Just for_iloop replaced :
image
Just for_iloop and einsum replaced:
image

I don't have a very solid reasoning behind why of it.


# 3. Pre-norm
layer_input = norm_fn(layer_input)
Expand All @@ -295,16 +294,13 @@ def __call__(
self.post_beta[...],
2.0,
)
post_out = jnp.einsum(
"bsd,bsk -> bskd",
layer_out,
post_mapping,
precision=self.matmul_precision,
)
# bsd,bsk -> bskd
post_out = jnp.expand_dims(layer_out, axis=2) * jnp.expand_dims(post_mapping, axis=3)

# 6. Residual mapping, res_out shape as [batch, seq, expansion_rate, emb]
res_mapping = self.res_mapping(norm_x)
res_out = jnp.einsum("bskd,bskm -> bsmd", x, res_mapping, precision=self.matmul_precision)
# bskd,bskm -> bsmd
res_out = jnp.sum(jnp.expand_dims(x, axis=3) * jnp.expand_dims(res_mapping, axis=4), axis=2)
return res_out + post_out, metadata


Expand Down
Loading