A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Bugfix] Allow skipping MoE in NVFP4 (fix for MTP) (#25987)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
Oct 6, 2025 at 16:16 UTC
2161efe9781cc0bee2f60342dffd0c1f7f0f2b57
5 files changed
+18
-5
vllm/model_executor/layers/fused_moe/layer.py
+2
@@ -1194,6 +1194,8 @@ class FusedMoE(CustomOp):
1194
if quant_config is None
1195
else quant_config.get_quant_method(self, prefix)
1196
)
1197
+ if quant_method is None:
1198
+ quant_method = UnquantizedFusedMoEMethod(moe)
1199
1200
assert quant_method is not None
1201
assert isinstance(quant_method, FusedMoEMethodBase)
vllm/model_executor/layers/quantization/modelopt.py
+4
-1
@@ -884,8 +884,9 @@ class ModelOptNvFp4Config(QuantizationConfig):
884
) -> Optional["QuantizeMethodBase"]:
885
from vllm.attention.layer import Attention # Avoid circular import
886
887
+ skip_layer = self.is_layer_excluded(prefix)
888
if isinstance(layer, LinearBase):
888
- if self.is_layer_excluded(prefix):
889
+ if skip_layer:
890
return UnquantizedLinearMethod()
891
# Check if this is a vision model layer that should not be quantized
892
if "vision_tower" in prefix or "vision_model" in prefix:
@@ -894,6 +895,8 @@ class ModelOptNvFp4Config(QuantizationConfig):
895
elif isinstance(layer, Attention):
896
return ModelOptFp8KVCacheMethod(self)
897
elif isinstance(layer, FusedMoE):
898
+ if skip_layer:
899
+ return None
900
return ModelOptNvFp4FusedMoE(self, layer.moe_config, layer)
901
return None
902
vllm/model_executor/models/deepseek_eagle.py
+1
@@ -55,6 +55,7 @@ class DeepseekV2Model(nn.Module):
55
DeepseekV2DecoderLayer(
56
vllm_config,
57
prefix=maybe_prefix(prefix, f"layers.{i + start_layer_id}"),
58
+ config=self.config,
59
)
60
for i in range(self.config.num_hidden_layers)
61
]
vllm/model_executor/models/deepseek_mtp.py
+7
-2
@@ -48,7 +48,8 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
48
def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
49
super().__init__()
50
51
- config = vllm_config.model_config.hf_config
51
+ config = vllm_config.speculative_config.draft_model_config.hf_config
52
+ self.config = config
53
quant_config = vllm_config.quant_config
54
55
self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -66,11 +67,15 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
67
)
68
else:
69
topk_indices_buffer = None
70
+
71
self.shared_head = SharedHead(
72
config=config, prefix=prefix, quant_config=quant_config
73
)
74
self.mtp_block = DeepseekV2DecoderLayer(
73
- vllm_config, prefix, topk_indices_buffer
75
+ vllm_config,
76
+ prefix,
77
+ config=self.config,
78
+ topk_indices_buffer=topk_indices_buffer,
79
)
80
81
def forward(
vllm/model_executor/models/deepseek_v2.py
+4
-2
@@ -1055,11 +1055,13 @@ class DeepseekV2DecoderLayer(nn.Module):
1055
self,
1056
vllm_config: VllmConfig,
1057
prefix: str,
1058
+ config: Optional[DeepseekV2Config] = None,
1059
topk_indices_buffer: Optional[torch.Tensor] = None,
1060
) -> None:
1061
super().__init__()
1062
1062
- config = vllm_config.model_config.hf_config
1063
+ if config is None:
1064
+ config = vllm_config.model_config.hf_config
1065
model_config = vllm_config.model_config
1066
cache_config = vllm_config.cache_config
1067
quant_config = vllm_config.quant_config
@@ -1200,7 +1202,7 @@ class DeepseekV2Model(nn.Module):
1202
self.start_layer, self.end_layer, self.layers = make_layers(
1203
config.num_hidden_layers,
1204
lambda prefix: DeepseekV2DecoderLayer(
1203
- vllm_config, prefix, topk_indices_buffer
1205
+ vllm_config, prefix, topk_indices_buffer=topk_indices_buffer
1206
),
1207
prefix=f"{prefix}.layers",
1208
)