From dc56df92ed678cd84fac2c6277c2fbb88d77c1ac Mon Sep 17 00:00:00 2001
From: Overwatch Agent <overwatch-agent@google.com>
Date: Fri, 7 Aug 2026 15:12:39 +0000
Subject: [PATCH] Fix DeepSeek4HyperHead instantiation error by using
 mhc_reduce

---
 src/maxtext/layers/mhc.py          | 54 ------------------------------
 src/maxtext/layers/nnx_decoders.py | 14 ++------
 2 files changed, 2 insertions(+), 66 deletions(-)

diff --git a/src/maxtext/layers/mhc.py b/src/maxtext/layers/mhc.py
index 80870aca46..93b172dfa1 100644
--- a/src/maxtext/layers/mhc.py
+++ b/src/maxtext/layers/mhc.py
@@ -313,58 +313,4 @@ def __call__(
     return res_out + post_out, metadata
 
 
-class DeepSeek4HyperHead(nnx.Module):
-  """DeepSeek V4 Hyper Head for collapsing hyper-connection streams."""
 
-  def __init__(self, config: Config, rngs: nnx.Rngs):
-    self.config = config
-    self.hc_mult = config.mhc_expansion_rate
-    self.eps = getattr(config, "hc_eps", 1e-6)
-    self.dtype = config.dtype
-    self.weight_dtype = config.weight_dtype
-    self.matmul_precision = jax.lax.Precision(config.matmul_precision)
-
-    self.input_norm = RMSNorm(
-        num_features=self.hc_mult * config.emb_dim,
-        dtype=self.dtype,
-        weight_dtype=self.weight_dtype,
-        kernel_axes=("norm",),
-        epsilon=config.normalization_layer_epsilon,
-        with_scale=False,
-        rngs=rngs,
-    )
-
-    scale_init = nd_dense_init(1.0, "fan_in", "normal")
-    self.hc_fn = nnx.Param(
-        scale_init(
-            rngs.params(),
-            (self.hc_mult * config.emb_dim, self.hc_mult),
-            self.weight_dtype,
-            in_axis=0,
-            out_axis=1,
-        ),
-        out_sharding=(None, None),
-    )
-
-    self.hc_base = nnx.Param(
-        default_bias_init(rngs.params(), (self.hc_mult,), self.weight_dtype),
-        out_sharding=(None,),
-    )
-    self.hc_scale = nnx.Param(
-        default_scalar_init(rngs.params(), (1,), self.weight_dtype),
-        out_sharding=(None,),
-    )
-
-  def __call__(self, x: Array) -> Array:
-    b, s, k, d = x.shape
-    flat = jnp.reshape(x, (b, s, k * d))
-    flat = self.input_norm(flat)
-
-    hc_fn = jnp.asarray(self.hc_fn[...], self.dtype)
-    hc_base = jnp.asarray(self.hc_base[...], self.dtype)
-    hc_scale = jnp.asarray(self.hc_scale[...], self.dtype)
-
-    mixes = jnp.einsum("bsm,mn -> bsn", flat, hc_fn, precision=self.matmul_precision)
-    pre = jax.nn.sigmoid(mixes * hc_scale + hc_base) + self.eps
-
-    return jnp.sum(jnp.expand_dims(pre, axis=-1) * x, axis=2).astype(self.dtype)
diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py
index 895ea27c14..3ce733e69b 100644
--- a/src/maxtext/layers/nnx_decoders.py
+++ b/src/maxtext/layers/nnx_decoders.py
@@ -438,13 +438,6 @@ def __init__(
     self.is_gemma4 = self.config.decoder_block == DecoderBlockType.GEMMA4
     self.is_gemma4_small = self.config.decoder_block == DecoderBlockType.GEMMA4_SMALL
 
-    if config.mhc_expansion_rate > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4:
-      self.hc_head = mhc.DeepSeek4HyperHead(
-          config=config,
-          mesh=self.mesh,
-          rngs=self.rngs,
-      )
-
     self._init_decoder_layers(decoder_block_classes, rngs, mesh)
 
   def _init_decoder_layers(self, decoder_block_classes, rngs, mesh):
@@ -1976,11 +1969,8 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in):
 
     # After the final transformer layer, `y` holds the raw, un-normalized hidden state.
     if getattr(cfg, "mhc_expansion_rate", 1) > 1:
-      if cfg.decoder_block == DecoderBlockType.DEEPSEEK4:
-        hidden_state = self.hc_head(y)
-      else:
-        # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim)
-        hidden_state = mhc_reduce(y)
+      # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim)
+      hidden_state = mhc_reduce(y)
     else:
       hidden_state = y