Fix DeepSeek4HyperHead instantiation error by using mhc_reduce

[olufiyin19]

Fix DeepSeek4HyperHead instantiation error by using mhc_reduce

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request removes the DeepSeek4HyperHead class from mhc.py and simplifies the decoder logic in nnx_decoders.py to use mhc_reduce for DEEPSEEK4 blocks. However, the reviewer points out that this removal will break the Linen-based decoder path in decoders.py which still references the removed hyper-head, and suggests updating it to prevent runtime errors.

I am having trouble creating individual review comments. Click here to see my feedback.

src/maxtext/layers/mhc.py (316-317)

Removing DeepSeek4HyperHead from mhc.py will break the Linen-based decoder path in src/maxtext/layers/decoders.py. Specifically, line 1280 in decoders.py still attempts to instantiate mhc.DeepSeek4HyperHeadToLinen:

    if cfg.mhc_expansion_rate > 1:
      if cfg.decoder_block == DecoderBlockType.DEEPSEEK4:
        hidden_state = mhc.DeepSeek4HyperHeadToLinen(
            config=cfg,
            mesh=mesh,
            name="hc_head",
        )(y)
      else:
        # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim)
        hidden_state = mhc_reduce(y)

Since DeepSeek4HyperHead is removed, this will raise an AttributeError at runtime. Please update src/maxtext/layers/decoders.py to also use mhc_reduce(y) for DecoderBlockType.DEEPSEEK4, similar to the changes made in nnx_decoders.py.