Skip to content

perf(deepseek_v4): use mx.fast.rms_norm in HyperConnection and HyperHead #16

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
0xClandestine wants to merge 1 commit into
base: pc/add-deepseekv4flash-model
Choose a base branch
from
Open

perf(deepseek_v4): use mx.fast.rms_norm in HyperConnection and HyperHead #16

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
22 changes: 5 additions & 17 deletions mlx_lm/models/deepseek_v4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -648,16 +648,10 @@ def _hc_expand_op(
return y.astype(block_out.dtype)


@mx.compile
def _rms_rsqrt(flat: mx.array, eps: float) -> mx.array:
return mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + eps)


@mx.compile
def _hc_mixes(flat: mx.array, fn_T: mx.array, norm_eps: float) -> mx.array:
"""Fused RMS-rsqrt + matmul + scale into single compiled graph."""
rsqrt = mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + norm_eps)
return (flat @ fn_T) * rsqrt
"""Fused RMS-norm + matmul: rms_norm(flat) @ fn_T."""
return mx.fast.rms_norm(flat, None, eps=norm_eps) @ fn_T


class HyperConnection(nn.Module):
Expand All @@ -679,11 +673,7 @@ def compute_weights(self, x: mx.array):
flat = x.reshape(B, L, H * D).astype(mx.float32)
if self._fn_T is None:
self._fn_T = self.fn.T
if self.training:
rsqrt = _rms_rsqrt(flat, self.norm_eps)
mixes = (flat @ self._fn_T) * rsqrt
else:
mixes = _hc_mixes(flat, self._fn_T, self.norm_eps)
mixes = _hc_mixes(flat, self._fn_T, self.norm_eps)
split_sinkhorn = _hc_split_sinkhorn_ops if self.training else hc_split_sinkhorn
return split_sinkhorn(
mixes,
Expand Down Expand Up @@ -753,8 +743,7 @@ def _hyper_head_op(
"""Fused HyperHead: RMS-rsqrt + matmul + sigmoid + weighted sum."""
B, L, H, D = x.shape
flat = x.reshape(B, L, H * D).astype(mx.float32)
rsqrt = mx.rsqrt((flat * flat).mean(axis=-1, keepdims=True) + norm_eps)
mixes = (flat @ fn.T) * rsqrt
mixes = mx.fast.rms_norm(flat, None, eps=norm_eps) @ fn.T
pre = mx.sigmoid(mixes * scale[0] + base) + hc_eps
return (pre[..., None] * x.astype(mx.float32)).sum(axis=2).astype(x.dtype)

Expand All @@ -778,8 +767,7 @@ def __call__(self, x: mx.array):
)
B, L, H, D = x.shape
flat = x.reshape(B, L, H * D).astype(mx.float32)
rsqrt = _rms_rsqrt(flat, self.norm_eps)
mixes = (flat @ self.fn.T) * rsqrt
mixes = mx.fast.rms_norm(flat, None, eps=self.norm_eps) @ self.fn.T
pre = mx.sigmoid(mixes * self.scale[0] + self.base) + self.hc_eps
return (pre[..., None] * x.astype(mx.float32)).sum(axis=2).astype(x.dtype)

Expand Down