Fix NVFP4 weights mapper: add prefix mappings, fix substr order

- Add orig_to_new_prefix mappings (layers→model.layers, embed_tokens→model.embed_tokens, etc.)
  AutoWeightsLoader strips the model. prefix before the mapper runs, so these are required
- Move .self_attn.compressor. → .attn.mla_attn.compressor. before .self_attn. → .attn.
  in substr_renames so compressor keys get the mla_attn prefix before the general rename
- Remove suffix renames (head.weight→lm_head.weight, embed.weight→embed_tokens.weight)
  that were causing double-mapping since the NVFP4 checkpoint already uses lm_head/embed_tokens
- Add unit test: tests/test_nvfp4_mapper.py (39 cases, no vLLM/CUDA needed)

tests/test_nvfp4_mapper.py

@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+"""Unit test for the NVFP4 weights mapper.
+
+Validates that checkpoint key names from our ModelOpt-quantized
+DeepSeek-V4-Pro checkpoint are correctly mapped to vLLM model
+parameter names.
+
+This can run WITHOUT vLLM or CUDA — it only tests the mapper logic.
+"""
+
+import re
+import sys
+from typing import Optional
+
+
+class WeightsMapper:
+    """Simplified WeightsMapper for testing."""
+
+    def __init__(
+        self,
+        orig_to_new_prefix: Optional[dict] = None,
+        orig_to_new_regex: Optional[dict] = None,
+        orig_to_new_suffix: Optional[dict] = None,
+        orig_to_new_substr: Optional[dict] = None,
+    ):
+        self.prefix_map = orig_to_new_prefix or {}
+        self.regex_map = orig_to_new_regex or {}
+        self.suffix_map = orig_to_new_suffix or {}
+        self.substr_map = orig_to_new_substr or {}
+
+    def map_name(self, name: str) -> str:
+        # 1. Prefix
+        for old, new in self.prefix_map.items():
+            if name.startswith(old):
+                name = new + name[len(old):]
+                break
+
+        # 2. Regex
+        for pattern, replacement in self.regex_map.items():
+            name = pattern.sub(replacement, name)
+
+        # 3. Suffix
+        for old, new in self.suffix_map.items():
+            if name.endswith(old):
+                name = name[: -len(old)] + new
+                break
+
+        # 4. Substr (ordered dict — specific before general)
+        for old, new in self.substr_map.items():
+            if old in name:
+                name = name.replace(old, new, 1)
+
+        return name
+
+
+def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper:
+    """Exact copy of the mapper from deepseek_v4.py."""
+    expert_rename_regex = {
+        re.compile(r"(\.experts\.\d+\.)gate_proj\."): r"\1w1.",
+        re.compile(r"(\.experts\.\d+\.)up_proj\."): r"\1w3.",
+        re.compile(r"(\.experts\.\d+\.)down_proj\."): r"\1w2.",
+    }
+
+    suffix_renames = {}
+
+    substr_renames = {
+        ".self_attn.compressor.": ".attn.mla_attn.compressor.",
+        ".self_attn.q_a_proj.": ".attn.wq_a.",
+        ".self_attn.kv_proj.": ".attn.wkv.",
+        ".self_attn.q_b_proj.": ".attn.wq_b.",
+        ".self_attn.o_a_proj.": ".attn.wo_a.",
+        ".self_attn.o_b_proj.": ".attn.wo_b.",
+        ".self_attn.q_a_norm.": ".attn.q_a_norm.",
+        ".self_attn.kv_norm.": ".attn.kv_norm.",
+        ".self_attn.sinks": ".attn.sinks",
+        ".mlp.shared_experts.gate_proj.": ".ffn.shared_experts.w1.",
+        ".mlp.shared_experts.up_proj.": ".ffn.shared_experts.w3.",
+        ".mlp.shared_experts.down_proj.": ".ffn.shared_experts.down_proj.",
+        ".mlp.": ".ffn.",
+        ".self_attn.": ".attn.",
+    }
+
+    return WeightsMapper(
+        orig_to_new_prefix={
+            "layers.": "model.layers.",
+            "embed_tokens.": "model.embed_tokens.",
+            "norm.": "model.norm.",
+            "hc_head": "model.hc_head",
+            "mtp.": "model.mtp.",
+        },
+        orig_to_new_regex=expert_rename_regex,
+        orig_to_new_suffix=suffix_renames,
+        orig_to_new_substr=substr_renames,
+    )
+
+
+TEST_CASES = [
+    # Embedding & top-level
+    ("embed_tokens.weight", "model.embed_tokens.weight"),
+    ("norm.weight", "model.norm.weight"),
+    ("hc_head.hc_fn", "model.hc_head.hc_fn"),
+    ("hc_head.hc_base", "model.hc_head.hc_base"),
+    ("hc_head.hc_scale", "model.hc_head.hc_scale"),
+    ("lm_head.weight", "lm_head.weight"),
+
+    # Attention — self_attn → attn
+    ("layers.0.self_attn.q_a_proj.weight", "model.layers.0.attn.wq_a.weight"),
+    ("layers.0.self_attn.q_a_proj.input_scale", "model.layers.0.attn.wq_a.input_scale"),
+    ("layers.0.self_attn.kv_proj.weight", "model.layers.0.attn.wkv.weight"),
+    ("layers.0.self_attn.kv_proj.input_scale", "model.layers.0.attn.wkv.input_scale"),
+    ("layers.0.self_attn.q_b_proj.weight", "model.layers.0.attn.wq_b.weight"),
+    ("layers.0.self_attn.o_a_proj.weight", "model.layers.0.attn.wo_a.weight"),
+    ("layers.0.self_attn.o_b_proj.weight", "model.layers.0.attn.wo_b.weight"),
+    ("layers.0.self_attn.o_b_proj.input_scale", "model.layers.0.attn.wo_b.input_scale"),
+    ("layers.0.self_attn.q_a_norm.weight", "model.layers.0.attn.q_a_norm.weight"),
+    ("layers.0.self_attn.kv_norm.weight", "model.layers.0.attn.kv_norm.weight"),
+    ("layers.0.self_attn.sinks", "model.layers.0.attn.sinks"),
+
+    # Attention — compressor (inside self_attn)
+    ("layers.0.self_attn.compressor.kv_proj.weight", "model.layers.0.attn.mla_attn.compressor.kv_proj.weight"),
+    ("layers.0.self_attn.compressor.kv_proj.input_scale", "model.layers.0.attn.mla_attn.compressor.kv_proj.input_scale"),
+    ("layers.0.self_attn.compressor.gate_proj.weight", "model.layers.0.attn.mla_attn.compressor.gate_proj.weight"),
+    ("layers.0.self_attn.compressor.gate_proj.input_scale", "model.layers.0.attn.mla_attn.compressor.gate_proj.input_scale"),
+    ("layers.0.self_attn.compressor.kv_norm.weight", "model.layers.0.attn.mla_attn.compressor.kv_norm.weight"),
+    ("layers.0.self_attn.compressor.position_bias", "model.layers.0.attn.mla_attn.compressor.position_bias"),
+
+    # MoE gate
+    ("layers.0.mlp.gate.tid2eid", "model.layers.0.ffn.gate.tid2eid"),
+    ("layers.0.mlp.gate.weight", "model.layers.0.ffn.gate.weight"),
+
+    # Expert weights — gate_proj → w1, up_proj → w3, down_proj → w2
+    ("layers.0.mlp.experts.0.gate_proj.weight", "model.layers.0.ffn.experts.0.w1.weight"),
+    ("layers.0.mlp.experts.0.up_proj.weight", "model.layers.0.ffn.experts.0.w3.weight"),
+    ("layers.0.mlp.experts.0.down_proj.weight", "model.layers.0.ffn.experts.0.w2.weight"),
+    ("layers.0.mlp.experts.0.gate_proj.input_scale", "model.layers.0.ffn.experts.0.w1.input_scale"),
+    ("layers.0.mlp.experts.0.gate_proj.weight_scale", "model.layers.0.ffn.experts.0.w1.weight_scale"),
+    ("layers.0.mlp.experts.0.gate_proj.weight_scale_2", "model.layers.0.ffn.experts.0.w1.weight_scale_2"),
+    ("layers.0.mlp.experts.255.down_proj.weight", "model.layers.0.ffn.experts.255.w2.weight"),
+
+    # Shared experts — gate_proj → w1, up_proj → w3, down_proj stays
+    ("layers.0.mlp.shared_experts.gate_proj.weight", "model.layers.0.ffn.shared_experts.w1.weight"),
+    ("layers.0.mlp.shared_experts.up_proj.weight", "model.layers.0.ffn.shared_experts.w3.weight"),
+    ("layers.0.mlp.shared_experts.down_proj.weight", "model.layers.0.ffn.shared_experts.down_proj.weight"),
+    ("layers.0.mlp.shared_experts.gate_proj.input_scale", "model.layers.0.ffn.shared_experts.w1.input_scale"),
+    ("layers.0.mlp.shared_experts.down_proj.weight_scale", "model.layers.0.ffn.shared_experts.down_proj.weight_scale"),
+    ("layers.0.mlp.shared_experts.down_proj.weight_scale_2", "model.layers.0.ffn.shared_experts.down_proj.weight_scale_2"),
+
+    # Layer norm
+    ("layers.0.post_attention_layernorm.weight", "model.layers.0.post_attention_layernorm.weight"),
+]
+
+
+def main():
+    mapper = _make_deepseek_v4_nvfp4_weights_mapper()
+    passed = 0
+    failed = 0
+
+    for ckpt_key, expected in TEST_CASES:
+        result = mapper.map_name(ckpt_key)
+        if result == expected:
+            passed += 1
+        else:
+            failed += 1
+            print(f"FAIL: {ckpt_key}")
+            print(f"  expected: {expected}")
+            print(f"  got:      {result}")
+
+    print(f"\n{passed} passed, {failed} failed")
+    return 0 if failed == 0 else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

vllm/patches/deepseek_v4.py

@@ -1648,12 +1648,14 @@ def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper:
     }
 
     suffix_renames = {
-        "head.weight": "lm_head.weight",
-        "embed.weight": "embed_tokens.weight",
+        # The NVFP4 checkpoint already uses lm_head / embed_tokens directly,
+        # no suffix renames needed (unlike the MXFP4 upstream format).
     }
 
     # NOTE: specific renames MUST come before general ones (applied in order)
     substr_renames = {
+        # Attention compressor (MUST come before .self_attn. → .attn.)
+        ".self_attn.compressor.": ".attn.mla_attn.compressor.",
         # Attention projections (specific before .self_attn. → .attn.)
         ".self_attn.q_a_proj.": ".attn.wq_a.",
         ".self_attn.kv_proj.": ".attn.wkv.",
@@ -1663,7 +1665,6 @@ def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper:
         ".self_attn.q_a_norm.": ".attn.q_a_norm.",
         ".self_attn.kv_norm.": ".attn.kv_norm.",
         ".self_attn.sinks": ".attn.sinks",
-        ".attn.compressor.": ".attn.mla_attn.compressor.",
         # Shared expert projections (specific before .mlp. → .ffn.)
         ".mlp.shared_experts.gate_proj.": ".ffn.shared_experts.w1.",
         ".mlp.shared_experts.up_proj.": ".ffn.shared_experts.w3.",
@@ -1675,7 +1676,11 @@ def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper:
 
     return WeightsMapper(
         orig_to_new_prefix={
-            # Keep model. prefix — model params use it (e.g. model.embed_tokens.weight)
+            "layers.": "model.layers.",
+            "embed_tokens.": "model.embed_tokens.",
+            "norm.": "model.norm.",
+            "hc_head": "model.hc_head",
+            "mtp.": "model.mtp.",
         },
         orig_to_new_regex=expert_rename_regex,
         orig_to_new_suffix=suffix_renames,