Fix DeepSeek4 shape mismatch in mhc.py and moe.py

src/maxtext/layers/mhc.py

@@ -313,22 +313,48 @@ def __call__(
     return res_out + post_out, metadata
 
 
-from maxtext.layers import linears
+class DeepSeek4HyperHead(nnx.Module):
+  """DeepSeek4 Hyper Head."""
 
+  def __init__(
+      self,
+      config: Config,
+      mesh: Mesh,
+      rngs: nnx.Rngs,
+  ):
+    self.config = config
+    self.mesh = mesh
+    self.rngs = rngs
+    self.dtype = self.config.dtype
+    self.weight_dtype = self.config.weight_dtype
+    self.k = config.mhc_expansion_rate
+    self.dim = config.emb_dim
 
-class DeepSeek4HyperHead(linears.DenseGeneral):
-  """DeepSeek V4 HyperHead for projecting expanded hidden states."""
-
-  def __init__(self, config: Config, mesh: Mesh, rngs: nnx.Rngs):
-    super().__init__(
-        in_features_shape=config.mhc_expansion_rate * config.emb_dim,
-        out_features_shape=config.emb_dim,
-        weight_dtype=config.weight_dtype,
-        kernel_axes=("mlp", "embed"),
-        rngs=rngs,
+    self.hc_base = nnx.Param(
+        jnp.zeros((self.dim,), dtype=self.weight_dtype),
+        out_sharding=(None,),
+    )
+    self.hc_fn = nnx.Param(
+        jnp.zeros((self.dim,), dtype=self.weight_dtype),
+        out_sharding=(None,),
+    )
+    self.hc_scale = nnx.Param(
+        jnp.zeros((self.dim,), dtype=self.weight_dtype),
+        out_sharding=(None,),
+    )
+    self.kernel = nnx.Param(
+        jnp.zeros((self.k, self.dim), dtype=self.weight_dtype),
+        out_sharding=(None, None),
     )
 
-  def __call__(self, x):
-    b, l, k, d = x.shape
-    x = jnp.reshape(x, (b, l, k * d))
-    return super().__call__(x)
+  def __call__(self, x: Array) -> Array:
+    kernel = jnp.asarray(self.kernel[...], self.dtype)
+    hc_base = jnp.asarray(self.hc_base[...], self.dtype)
+    hc_fn = jnp.asarray(self.hc_fn[...], self.dtype)
+    hc_scale = jnp.asarray(self.hc_scale[...], self.dtype)
+
+    x = x * kernel
+    x = jnp.sum(x, axis=2)
+    x = x * hc_scale + hc_base + hc_fn
+    return x
+    return jnp.sum(x, axis=2)

src/maxtext/layers/moe.py

@@ -420,22 +420,23 @@ def __init__(
     self.rngs = rngs
     self.is_hash_routing = is_hash_routing
 
+    # DeepSeek V4 Hash Routing
     # DeepSeek V4 Hash Routing
     if self.is_hash_routing:
       # Token-ID to Expert-ID lookup table for static routing
       # Must be stored as float32 because MaxText passes the entire variable tree
       # through jax.value_and_grad, which strictly requires all leaves to be inexact types
       # (even if they receive no gradients). We cast to int32 dynamically during routing.
+      vocab_size = getattr(self.config, "original_vocab_size", self.config.vocab_size)
       self.tid2eid = Tid2EidVar(
           jnp.zeros(
-              (self.config.vocab_size, self.num_experts_per_tok),
+              (vocab_size, self.num_experts_per_tok),
               dtype=jnp.float32,
           ),
           out_sharding=None,  # Replicated across shards for local lookup
       )
     else:
       self.tid2eid = None
-
     self.moe_expert_input_dim = (
         self.config.emb_dim if self.config.moe_expert_input_dim <= 0 else self.config.moe_expert_input_dim
     )

src/maxtext/layers/nnx_decoders.py

@@ -1969,6 +1969,12 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in):
 
           if deepstack_visual_embeds is not None and lyr < len(deepstack_visual_embeds):
             visual_embeds = deepstack_visual_embeds[lyr]
+            if bidirectional_mask is not None and visual_embeds is not None:
+              y = deepstack_process(y, bidirectional_mask, visual_embeds)
+
+    assert isinstance(y, jax.Array)
+
+    # 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)
@@ -1977,14 +1983,6 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in):
         hidden_state = mhc_reduce(y)
     else:
       hidden_state = y
-    assert isinstance(y, jax.Array)
-
-    # After the final transformer layer, `y` holds the raw, un-normalized hidden state.
-    if getattr(cfg, "mhc_expansion_rate", 1) > 1:
-      # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim)
-      hidden_state = mhc_reduce(y)
-    else:
-      hidden_state = y
 
     # When invoking from vLLM with RPA attention, logit computation is deferred to a later stage.
     if cfg.attention in ("vllm_rpa", "vllm_batched_rpa"):