Address PR feedback for DeepSeek-V4 integration

Author: parambole

src/maxtext/configs/models/deepseek4-284b.yml

@@ -36,9 +36,10 @@ indexer_head_dim: 128
 indexer_n_heads: 64
 indexer_topk: 512
 
-# Note: Layers (0, 1, 2) are prefix layers.
+# Note: Layers (0, 1, 2) are prefix layers as `first_num_hash_layers=3`.
 # The 44th layer (MTP module with compress_ratio=0) has been explicitly dropped for now.
 # This leaves exactly 43 layers: 3 prefix [0,0,4] + 40 scanned.
+# `compress_ratio=0` uses sliding window attention. In this case, layer (0, 1).
 compress_ratios: [0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4]
 
 # --- MoE configuration ---
@@ -57,7 +58,6 @@ o_lora_rank: 1024
 sliding_window_size: 128
 
 # --- RoPE ---
-
 rope_type: "default"
 rope_max_timescale: 10000                  # Main RoPE theta
 compressed_rope_max_timescale: 160000      # Compressed RoPE theta

src/maxtext/layers/attention_compressed.py

@@ -697,6 +697,9 @@ def __init__(
     self.q_lora_rank = q_lora_rank
     self.compress_ratio = compress_ratio
 
+    if self.compress_ratio == 0:
+      attention_type = AttentionType.LOCAL_SLIDING
+
     super().__init__(
         config=config,
         num_query_heads=num_query_heads,

src/maxtext/layers/decoders.py

@@ -1238,6 +1238,7 @@ def __call__(
                 previous_chunk=previous_chunk,
                 slot=slot,
                 kv_cache=kv_cache,
+                decoder_input_tokens=decoder_input_tokens,
                 attention_metadata=attention_metadata,
                 **layer_call_kwargs,
             )
@@ -1481,7 +1482,12 @@ def _apply_deepseek4_scanned_blocks(
     }
 
     # 1. Prefix Unrolling
-    # These layers use Hash Routing.
+    # Prefix layers are unrolled (unscanned) for two architectural reasons:
+    #   1. Heterogeneous Attention: JAX nn.scan requires identical computation graphs, but the first few layers
+    #      use different attention configurations (e.g., DeepSeek-V4 uses compress_ratios [0, 0, 4] for layers 0, 1, 2).
+    #   2. Static Hash Routing: The first `first_num_hash_layers` (which is 3 for DeepSeek-V4) use deterministic
+    #      token-to-expert Hash Routing instead of learned top-k routing.
+    # Therefore, these prefix layers are instantiated individually before we scan the remaining uniform blocks.
     num_hash_layers = cfg.first_num_hash_layers
     for layer_idx in range(num_hash_layers):
       prefix_layer = deepseek4.DeepSeek4LayerToLinen(

src/maxtext/models/deepseek4.py

@@ -93,12 +93,6 @@ def __init__(
     if compress_ratio is None:
       compress_ratio = config.compress_ratios[layer_idx]
 
-    # Route to LOCAL_SLIDING if compression is disabled for this layer,
-    # otherwise default to the globally configured attention type (e.g., COMPRESSED).
-    layer_attention_type = (
-        AttentionType.LOCAL_SLIDING if compress_ratio == 0 else AttentionType(self.config.attention_type)
-    )
-
     self.self_attention = attention_compressed.CompressedAttention(
         config=self.config,
         compress_ratio=compress_ratio,
@@ -108,7 +102,7 @@ def __init__(
         max_target_length=self.config.max_target_length,
         max_prefill_predict_length=self.config.max_prefill_predict_length,
         attention_kernel=self.config.attention,
-        attention_type=layer_attention_type,
+        attention_type=AttentionType(self.config.attention_type),
         inputs_q_shape=self.dummy_inputs_shape,
         inputs_kv_shape=self.dummy_inputs_shape,
         mesh=self.mesh,

src/maxtext/utils/globals.py

@@ -75,7 +75,7 @@
     "deepseek2-16b": "deepseek-ai/DeepSeek-V2-Lite",
     "deepseek3-671b": "deepseek-ai/DeepSeek-V3",
     "deepseek3.2-671b": "deepseek-ai/DeepSeek-V3.2",
-    "deepseek4": "deepseek-ai/DeepSeek-V4-Flash",
+    "deepseek4-284b": "deepseek-ai/DeepSeek-V4-Flash",
     "gpt-oss-20b": "openai/gpt-oss-20b",
     "gpt-oss-120b": "openai/gpt-oss-120b",
     "qwen3-omni-30b-a3b": "Qwen/Qwen3-Omni-30B-A3B-Instruct",

tests/unit/train_compile_test.py

@@ -804,9 +804,14 @@ def test_deepseek32(self):
         )
     )
 
-  def test_deepseek4(self):
+  @parameterized.named_parameters(
+      {"testcase_name": "scanned", "scan_layers": "true"},
+      {"testcase_name": "unscanned", "scan_layers": "false"},
+  )
+  @pytest.mark.cpu_only
+  def test_deepseek4(self, scan_layers):
     # test deepseek4 compile
-    compiled_trainstep_file = "/tmp/test_deepseek4.pickle"
+    compiled_trainstep_file = f"/tmp/test_deepseek4_{scan_layers}.pickle"
     train_compile_main(
         (
             "",
@@ -818,6 +823,7 @@ def test_deepseek4(self):
             "model_name=deepseek4-284b",
             "per_device_batch_size=1",
             "max_target_length=1024",
+            f"scan_layers={scan_layers}",
             "attention=dot_product",
             "dtype=bfloat16",
             "weight_dtype=bfloat16",