[DeepSeek-V4] Implement MoE routing primitives

[parambole]

Description

Implement Mixture of Experts (MoE) routing primitives and infrastructural updates required for DeepSeek-V4 integration into MaxText.

• DeepSeek Routing Primitives: Integrated token-based hash routing logic and DeepSeek-specific Top-K routing directly into moe.py's get_topk function.
• Safe Layer Parameters: Updated RoutedMoE and RoutedAndSharedMoE to accept layer_idx and input_ids via optional keyword arguments. This allows DeepSeek to pipe these variables down into the hashing logic, while safely defaulting to 0 and None so existing callers are entirely unimpacted.
• Hardware Optimization Bypasses: Added boundary checks to bypass the TPU fused_moe_matmul (via the vllm_rpa block) specifically for hash-routed layers where it is unsupported.
• Activation Functions: Registered the sqrtsoftplus custom activation function in linears.py required by the DeepSeek V4 MoE router.
• Unit test suite (tests/unit/deepseek_v4_vs_reference_test.py): Validates MoE routing parity against PyTorch reference implementations.

Tests

Tested on CPU. All DeepSeek V4 and core MoE unit tests pass successfully.

pytest tests/unit/deepseek_v4_vs_reference_test.py tests/unit/moe_test.py
============================= test session starts ==============================
collecting ... 
collected 43 items
tests/unit/deepseek_v4_vs_reference_test.py ............                 [ 27%]
tests/unit/moe_test.py ...............................                   [100%]
================ 25 passed, 18 skipped, 123 warnings in 47.50s =================

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 61.53846% with 15 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/moe.py	64.86%	8 Missing and 5 partials ⚠️
src/maxtext/layers/linears.py	0.00%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @parambole, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[dipakg-lang]

Looks like this change wrongly has attention related changes instead of MOE routing related ones?

[parambole]

Looks like this change wrongly has attention related changes instead of MOE routing related ones?

Ah! That was just an artifact of the base branch I selected. I recently changed the PR's base branch to deepseek_v4_compressed_attention, which should isolate the diff to show only the MoE routing changes going forward. Let me know if the diff still looks noisy!

[RissyRan]

Thanks for the change! Please note that we have DS v2, and a few if condition may break the functionality.

[shuningjin]

Thanks for refactoring and reusing MoE logic! Further comments:

• Perhaps we can use is_hash_routing instead layer_idx to indicate hash router for clarify. Decoder can use idx to determine.
• Perhaps can use nnx.Variable instead of nnx.Param for register_buffer counterpart

[RissyRan]

Thanks for the change! Please add an assertion in types.py.

[shuningjin]

Great work!

[shuningjin]

Note sure why tests/integration/hlo_diff_test.py fails. Please double check.