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
[Model] MTP fallback to eager for DeepSeek v32 (#25982)
Signed-off-by: Lu Fang <fanglu@fb.com>
Lucia Fang committed
Sep 30, 2025 at 18:53 UTC
001e50c92c1557bb8c4afc3de0e2d8fc3bda348c
5 files changed
+32
-5
tests/v1/spec_decode/test_eagle.py
+10
-1
@@ -337,13 +337,19 @@ def test_load_model(mock_get_model, mock_get_layers, mock_get_pp_group, method,
337
"target_attn_1": mock.MagicMock(),
338
"target_attn_2": mock.MagicMock()
339
}
340
+ target_indx_layers: dict[str, mock.MagicMock] = {}
341
# Draft model has one extra attention layer compared to target model
342
all_attn_layers = {
343
**target_attn_layers, "draft_extra_attn": mock.MagicMock()
344
}
345
346
+ all_indx_layers: dict[str, mock.MagicMock] = {}
347
+
348
# Make mock_get_layers return different values for each call
346
- mock_get_layers.side_effect = [target_attn_layers, all_attn_layers]
349
+ mock_get_layers.side_effect = [
350
+ target_attn_layers, target_indx_layers, all_attn_layers,
351
+ all_indx_layers
352
+ ]
353
354
# Setup mock for pp group to return the appropriate value for world size
355
mock_pp_group = mock.MagicMock()
@@ -658,6 +664,9 @@ def test_propose_tree(spec_token_tree):
664
# Mock runner for attention metadata building.
665
proposer.runner = mock.MagicMock()
666
proposer.runner.attn_groups.append([mock.MagicMock()])
667
+ proposer.runner.attn_groups[0][0].metadata_builders = [
668
+ attn_metadata_builder
669
+ ]
670
proposer.runner.attn_groups[0][0].get_metadata_builder.return_value = \
671
attn_metadata_builder
672
proposer._get_attention_metadata_builder = mock.MagicMock(
tests/v1/spec_decode/test_mtp.py
+7
-1
@@ -63,7 +63,13 @@ def test_mtp_load_model_unified(mock_get_model, mock_get_layers,
63
64
target_attn_layers = {"target_attn_1": mock.MagicMock()}
65
all_attn_layers = {**target_attn_layers, "draft_attn_1": mock.MagicMock()}
66
- mock_get_layers.side_effect = [target_attn_layers, all_attn_layers]
66
+ target_indexer_layers: dict = {}
67
+ all_indexer_layers: dict = {}
68
+
69
+ mock_get_layers.side_effect = [
70
+ target_attn_layers, target_indexer_layers, all_attn_layers,
71
+ all_indexer_layers
72
+ ]
73
74
mock_pp_group = mock.MagicMock()
75
mock_pp_group.world_size = 1
vllm/config/speculative.py
+7
-1
@@ -41,7 +41,8 @@ MTP_MODEL_TYPES = ("deepseek_mtp", "mimo_mtp", "glm4_moe_mtp", "ernie_mtp",
41
@dataclass
42
class SpeculativeConfig:
43
"""Configuration for speculative decoding."""
44
-
44
+ enforce_eager: Optional[bool] = None
45
+ """Override the default enforce_eager from model_config"""
46
# General speculative decoding control
47
num_speculative_tokens: SkipValidation[int] = None # type: ignore
48
"""The number of speculative tokens, if provided. It will default to the
@@ -219,6 +220,11 @@ class SpeculativeConfig:
220
assert (
221
self.target_model_config
222
is not None), "target_model_config must be present for mtp"
223
+ if self.target_model_config.hf_text_config.model_type \
224
+ == "deepseek_v32":
225
+ # FIXME(luccafong): cudgraph with v32 MTP is not supported,
226
+ # remove this when the issue is fixed.
227
+ self.enforce_eager = True
228
# use the draft model from the same model:
229
self.model = self.target_model_config.model
230
# Align the quantization of draft model for cases such as
vllm/v1/attention/backends/mla/indexer.py
+1
-1
@@ -171,7 +171,7 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig):
171
172
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
173
cudagraph_support: ClassVar[AttentionCGSupport] = \
174
- AttentionCGSupport.UNIFORM_BATCH
174
+ AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
175
176
reorder_batch_threshold: int = 1
177
vllm/v1/spec_decode/eagle.py
+7
-1
@@ -50,6 +50,7 @@ class EagleProposer:
50
):
51
self.vllm_config = vllm_config
52
self.speculative_config = vllm_config.speculative_config
53
+ assert self.speculative_config is not None
54
self.draft_model_config = self.speculative_config.draft_model_config
55
self.method = self.speculative_config.method
56
@@ -74,11 +75,16 @@ class EagleProposer:
75
vllm_config.model_config)
76
77
self.attn_metadata_builder: Optional[AttentionMetadataBuilder] = None
78
+ self.draft_indexer_metadata_builder: Optional[
79
+ AttentionMetadataBuilder] = None
80
+ self.attn_layer_names: list[str] = []
81
+ self.indexer_layer_names: list[str] = []
82
83
self.use_cuda_graph = (not current_platform.is_xpu()
84
and self.vllm_config.compilation_config.level
85
== CompilationLevel.PIECEWISE and
81
- not self.vllm_config.model_config.enforce_eager)
86
+ not self.vllm_config.model_config.enforce_eager
87
+ and not self.speculative_config.enforce_eager)
88
self.cudagraph_batch_sizes = list(
89
reversed(self.vllm_config.compilation_config.
90
cudagraph_capture_sizes)) if self.use_cuda_graph else []