[DeepSeek-V4] Implement model integration, decoders, and configuration stack

[parambole]

Description

This PR introduces native architectural and routing support for the DeepSeek V4 model in MaxText.

Why & What: DeepSeek V4 introduces non-uniform architectural features that require explicit configuration unrolling. This PR solves the integration by implementing:

• Compressed Attention (CSA/HCA): Bypasses standard MLA instantiation and natively integrates DeepSeek V4's alternating CSA and HCA attention blocks.
• Hybrid Routing: Implements DeepSeek's transition from fixed Hash Routing (early layers) to learned Token Routing (later layers) natively within the MoE framework.
• Architectural Scanning: Unrolls the 44-layer configuration to properly handle the [0, 0] prefix compression ratios, the perfectly alternating [4, 128] scanned middle layers, and the [4, 0] suffix layers.

Tests

• Unit Tests: Verified mathematical parity against reference implementations using tests/unit/deepseek_v4_vs_reference_test.py.
• E2E Compilation: Successfully compiled the full DeepSeek V4 model on a simulated v5p-512 mesh to guarantee memory constraints and HLO generation.

Compile Command to Reproduce:

python3  -m  maxtext.trainers.pre_train.train_compile  src/maxtext/configs/base.yml
  base_output_directory=/tmp/maxtext_logs
  run_name=dsv4_v5p512_compile
  per_device_batch_size=1
  enable_checkpointing=false
  model_name=deepseek4
  compile_topology=v5p-512
  compile_topology_num_slices=1
  ici_fsdp_parallelism=-1
  steps=1
  max_target_length=4096
  async_checkpointing=false
  tokenizer_type=huggingface
  tokenizer_path=deepseek-ai/DeepSeek-V3
  attention=dot_product
  dtype=bfloat16
  weight_dtype=bfloat16
  megablox=False
  sparse_matmul=False
  dataset_type=synthetic
  scan_layers=true

Proof of Compilation:

Memory analysis: CompiledMemoryStats(generated_code_size_in_bytes=260855808, argument_size_in_bytes=18401962496, output_size_in_bytes=18401889280, alias_size_in_bytes=18401880576, temp_size_in_bytes=94892786400, host_generated_code_size_in_bytes=0, host_argument_size_in_bytes=0, host_output_size_in_bytes=0, host_alias_size_in_bytes=0, host_temp_size_in_bytes=0)

Checklist

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 91.75258% with 8 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/decoders.py	83.33%	1 Missing and 3 partials ⚠️
src/maxtext/layers/moe.py	71.42%	1 Missing and 1 partial ⚠️
src/maxtext/models/deepseek4.py	96.36%	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.

[shuningjin]

Thank you!

1. The unscanned seems have error. Could you fix and add it to train compile unit test as well?

python -m maxtext.trainers.pre_train.train_compile src/maxtext/configs/base.yml compile_topology=v5p-256 compile_topology_num_slices=1 model_name=deepseek4-284b per_device_batch_size=1 scan_layers=false run_name=test attention=dot_product

https://paste.googleplex.com/6636383660802048

2. The param structure seems different from other models. Could you make it more consistent?
   For the scanned structure, while it seems different from deepseek3, it is consistent with gemma4. I guess this is fine

---

python src/maxtext/checkpoint_conversion/inspect_checkpoint.py maxtext \
model_name=deepseek3-671b scan_layers=true

https://paste.googleplex.com/6550606788558848

python src/maxtext/checkpoint_conversion/inspect_checkpoint.py maxtext \
model_name=deepseek3-671b scan_layers=false

https://paste.googleplex.com/6331198317133824

python src/maxtext/checkpoint_conversion/inspect_checkpoint.py maxtext \
model_name=llama4-17b-16e scan_layers=true

https://paste.googleplex.com/6228521688104960

python src/maxtext/checkpoint_conversion/inspect_checkpoint.py maxtext \
model_name=gemma4-26b scan_layers=true

https://paste.googleplex.com/4838680605884416

python src/maxtext/checkpoint_conversion/inspect_checkpoint.py maxtext \
model_name=deepseek4-284b scan_layers=true

https://paste.googleplex.com/6297935070494720

python src/maxtext/checkpoint_conversion/inspect_checkpoint.py maxtext \
model_name=deepseek4-284b scan_layers=false

error: https://paste.googleplex.com/6194718466768896

[parambole]

python src/maxtext/checkpoint_conversion/inspect_checkpoint.py maxtext \
model_name=deepseek4-284b scan_layers=false

error: https://paste.googleplex.com/6194718466768896

@shuningjin Great catch and thank you for adding the tool. I have fixed the error and rerun the test

https://paste.googleplex.com/5780004561616896

[shuningjin]

Excellent job implementing such a novel and complex model! Really appreciate your hard work.

[NuojCheng]

Great work!!