biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix NVFP4 mapper: layer norms, hc params, indexer path, q_a_norm

Browse Source 
- input_layernorm → attn_norm, post_attention_layernorm → ffn_norm
- hc_head.fn/base/scale → hc_head_fn/base/scale
- attn_hc/ffn_hc → hc_attn/hc_ffn (dot to underscore)
- q_a_norm → q_norm, sinks → attn_sink
- Indexer params: self_attn.compressor.indexer → attn.indexer
  (not attn.mla_attn.compressor.indexer)
This commit is contained in:
biondizzle
2026-05-19 00:24:26 +00:00
parent 540e7ee8fc
commit 10c14ddb49
1 changed files with 34 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
vllm/patches/deepseek_v4.py
View File

@@ -1771,6 +1771,18 @@ def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper:
# NOTE: specific renames MUST come before general ones (applied in order)
substr_renames = {
# === Indexer params (MUST come before .self_attn.compressor.
# so that indexer keys are captured before the compressor prefix
# rewrite moves them under mla_attn.compressor) ===
# The checkpoint puts indexer under self_attn.compressor.indexer.*
# but the model has indexer at attn.indexer.* (sibling of compressor,
# NOT nested under it).
".self_attn.compressor.indexer.q_b_proj.": ".attn.indexer.wq_b.",
".self_attn.compressor.indexer.weights_proj.": ".attn.indexer.weights_proj.",
".self_attn.compressor.indexer.kv_norm.": ".attn.indexer.k_norm.",
".self_attn.compressor.indexer.kv_proj.": ".attn.indexer.compressor.wkv.",
".self_attn.compressor.indexer.gate_proj.": ".attn.indexer.compressor.wgate.",
".self_attn.compressor.indexer.position_bias": ".attn.indexer.compressor.ape",
# === Compressor (non-indexer) NVFP4 renames ===
# Checkpoint uses kv_proj/gate_proj, model uses wkv/wgate
# (for stacking into fused_wkv_wgate).
@@ -1778,29 +1790,19 @@ def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper:
"compressor.gate_proj.": "compressor.wgate.",
"compressor.kv_norm.": "compressor.norm.",
"compressor.position_bias": "compressor.ape",
# === Attention compressor (MUST come before .self_attn. → .attn.
# and before indexer renames so that .self_attn.compressor.indexer.
# becomes .attn.mla_attn.compressor.indexer.) ===
# === Attention compressor (MUST come after indexer renames
# so that remaining .self_attn.compressor. (non-indexer) keys
# become .attn.mla_attn.compressor.) ===
".self_attn.compressor.": ".attn.mla_attn.compressor.",
# === Indexer params (now under .attn.mla_attn.compressor.indexer.
# after the compressor rename above) ===
# Indexer's own params
"compressor.indexer.q_b_proj.": "indexer.wq_b.",
"compressor.indexer.weights_proj.": "indexer.weights_proj.",
"compressor.indexer.kv_norm.": "indexer.k_norm.",
# Indexer's compressor (compressor.indexer → indexer.compressor)
"compressor.indexer.kv_proj.": "indexer.compressor.wkv.",
"compressor.indexer.gate_proj.": "indexer.compressor.wgate.",
"compressor.indexer.position_bias": "indexer.compressor.ape",
# === Attention projections (specific before .self_attn. → .attn.) ===
".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.q_a_norm.": ".attn.q_norm.",
".self_attn.kv_norm.": ".attn.kv_norm.",
".self_attn.sinks": ".attn.sinks",
".self_attn.sinks": ".attn.attn_sink",
# 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.",
@@ -1808,6 +1810,23 @@ def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper:
# General renames
".mlp.": ".ffn.",
".self_attn.": ".attn.",
# Layer norms (checkpoint uses input_layernorm / post_attention_layernorm,
# model uses attn_norm / ffn_norm)
"input_layernorm.": "attn_norm.",
"post_attention_layernorm.": "ffn_norm.",
# Per-layer HC params (checkpoint uses attn_hc / ffn_hc with dot,
# model uses hc_attn / hc_ffn with underscore)
".attn_hc.fn": ".hc_attn_fn",
".attn_hc.base": ".hc_attn_base",
".attn_hc.scale": ".hc_attn_scale",
".ffn_hc.fn": ".hc_ffn_fn",
".ffn_hc.base": ".hc_ffn_base",
".ffn_hc.scale": ".hc_ffn_scale",
# Top-level hc_head params (checkpoint uses hc_head.fn etc,
# model uses hc_head_fn etc)
"hc_head.fn": "hc_head_fn",
"hc_head.base": "hc_head_base",
"hc_head.scale": "hc_head_scale",
}
return WeightsMapper(