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 15, 21:00 UTC
[Model][Speculative Decoding] Expand DeepSeek MTP code to support k > n_predict (#13626)
Signed-off-by: Benjamin Chislett <benjamin.chislett@centml.ai>
Benjamin Chislett committed
Feb 27, 2025 at 18:28 UTC
9804145cac8b108725b2a431be6cd8a65d20ef9a
6 files changed
+49
-22
vllm/config.py
+5
-6
@@ -1978,13 +1978,12 @@ class SpeculativeConfig:
1978
if num_speculative_tokens is None:
1979
# Default to max value defined in draft model config.
1980
num_speculative_tokens = n_predict
1981
- elif num_speculative_tokens > n_predict:
1982
- # Verify provided value doesn't exceed the maximum
1983
- # supported by the draft model.
1981
+ elif num_speculative_tokens > n_predict and \
1982
+ num_speculative_tokens % n_predict != 0:
1983
+ # Ensure divisibility for MTP module reuse.
1984
raise ValueError(
1985
- "This speculative model supports a maximum of "
1986
- f"num_speculative_tokens={n_predict}, but "
1987
- f"{num_speculative_tokens=} was provided.")
1985
+ f"{num_speculative_tokens=} must be divisible by "
1986
+ f"{n_predict=}")
1987
1988
speculative_draft_tensor_parallel_size = \
1989
SpeculativeConfig._verify_and_get_draft_model_tensor_parallel_size(
vllm/model_executor/models/deepseek_mtp.py
+9
-5
@@ -87,7 +87,7 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
87
hidden_states=hidden_states,
88
residual=None)
89
hidden_states = residual + hidden_states
90
- return self.shared_head(hidden_states)
90
+ return hidden_states
91
92
93
class DeepSeekMultiTokenPredictor(nn.Module):
@@ -121,12 +121,13 @@ class DeepSeekMultiTokenPredictor(nn.Module):
121
inputs_embeds: Optional[torch.Tensor] = None,
122
spec_step_idx: int = 0,
123
) -> torch.Tensor:
124
- return self.layers[str(self.mtp_start_layer_idx + spec_step_idx)](
124
+ current_step_idx = (spec_step_idx % self.num_mtp_layers)
125
+ return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
126
input_ids,
127
positions,
128
previous_hidden_states,
129
inputs_embeds,
129
- spec_step_idx,
130
+ current_step_idx,
131
)
132
133
def compute_logits(
@@ -135,9 +136,12 @@ class DeepSeekMultiTokenPredictor(nn.Module):
136
sampling_metadata: SamplingMetadata,
137
spec_step_idx: int = 0,
138
) -> torch.Tensor:
138
- mtp_layer = self.layers[str(self.mtp_start_layer_idx + spec_step_idx)]
139
+ current_step_idx = (spec_step_idx % self.num_mtp_layers)
140
+ mtp_layer = self.layers[str(self.mtp_start_layer_idx +
141
+ current_step_idx)]
142
logits = self.logits_processor(mtp_layer.shared_head.head,
140
- hidden_states, sampling_metadata)
143
+ mtp_layer.shared_head(hidden_states),
144
+ sampling_metadata)
145
return logits
146
147
vllm/spec_decode/draft_model_runner.py
+4
-7
@@ -50,12 +50,6 @@ class TP1DraftModelRunner(ModelRunnerWrapperBase):
50
"""
51
52
def __init__(self, model_runner: ModelRunnerBase):
53
- if hasattr(
54
- model_runner,
55
- "return_hidden_states") and model_runner.return_hidden_states:
56
- raise ValueError(
57
- "return_hidden_states is not supported for TP1DraftModelRunner."
58
- )
53
super().__init__(model_runner)
54
55
self.indices_of_seq_with_bonus_tokens = None
@@ -153,7 +147,7 @@ class TP1DraftModelRunner(ModelRunnerWrapperBase):
147
return False
148
149
# TODO: Add support for other attn backends
156
- if self.attn_backend.get_name() not in ("FLASH_ATTN", "TRITON_MLA"):
150
+ if self.attn_backend.get_name() not in ("FLASH_ATTN", ):
151
return False
152
153
# TODO: Add support for LORA
@@ -307,6 +301,9 @@ class TP1DraftModelRunner(ModelRunnerWrapperBase):
301
)
302
outputs.append(output)
303
304
+ if self.return_hidden_states and is_fallback:
305
+ output.hidden_states = hidden_states
306
+
307
if model_input.attn_metadata.num_prefills == 0 \
308
and self.indices_of_seq_with_bonus_tokens is not None:
309
assert output.sampled_token_ids is not None
vllm/spec_decode/multi_step_worker.py
+17
@@ -96,12 +96,16 @@ class MultiStepWorker(ProposerWorkerBase, DelegateWorkerBase):
96
# TODO: Remove this branch once DraftModelRunner supports TP>1
97
# and other restrictions that are part of DraftModelRunner's
98
# supports_gpu_multi_step(..)
99
+ if expanded_request.previous_hidden_states is not None:
100
+ self.worker.model_runner.return_hidden_states = True
101
for _ in range(sample_len):
102
model_output: List[SamplerOutput] = self.worker.execute_model(
103
execute_model_req=expanded_request)
104
assert (len(model_output) == 1
105
), "composing multistep workers not supported"
106
model_output = model_output[0]
107
+ self._maybe_update_previous_hidden_states(
108
+ model_output, expanded_request)
109
110
self._append_new_tokens(
111
model_output, expanded_request.seq_group_metadata_list,
@@ -115,6 +119,19 @@ class MultiStepWorker(ProposerWorkerBase, DelegateWorkerBase):
119
model_outputs, indices_of_seq_with_bonus_tokens)
120
return filtered_model_outputs, True
121
122
+ @staticmethod
123
+ def _maybe_update_previous_hidden_states(
124
+ model_output: SamplerOutput,
125
+ expanded_request: ExecuteModelRequest) -> None:
126
+ """
127
+ Updates the previous hidden states in an expanded request
128
+ in-place with the hidden states from the model output.
129
+ """
130
+ if expanded_request.previous_hidden_states is not None:
131
+ expanded_request.previous_hidden_states = HiddenStates(
132
+ model_output.hidden_states,
133
+ expanded_request.seq_group_metadata_list)
134
+
135
@staticmethod
136
def _expand_execute_model_request(
137
execute_model_req: ExecuteModelRequest,
vllm/spec_decode/spec_decode_worker.py
+3
-3
@@ -184,8 +184,7 @@ class SpecDecodeWorker(LoRANotSupportedWorkerBase):
184
elif draft_model_config.hf_config.model_type == "medusa":
185
proposer_worker = MedusaWorker(**draft_worker_kwargs)
186
else:
187
- if draft_tp == 1 or draft_model_config.hf_config.model_type ==\
188
- "deepseek_mtp":
187
+ if draft_tp == 1:
188
if current_platform.is_cuda_alike():
189
draft_worker_kwargs[
190
"model_runner_cls"] = TP1DraftModelRunner
@@ -203,7 +202,8 @@ class SpecDecodeWorker(LoRANotSupportedWorkerBase):
202
203
proposer_worker = MultiStepWorker(**draft_worker_kwargs)
204
if draft_model_config.hf_config.model_type == "deepseek_mtp":
206
- num_spec_prefill_steps = num_speculative_tokens
205
+ num_spec_prefill_steps = \
206
+ draft_model_config.hf_config.n_predict
207
208
proposer_worker = SmallerTpProposerWorker.maybe_wrap_worker(
209
proposer_worker, draft_tp, target_tp)
vllm/worker/model_runner.py
+11
-1
@@ -1685,11 +1685,22 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
1685
# TODO(andoorve): We can remove this once all
1686
# virtual engines share the same kv cache.
1687
virtual_engine = model_input.virtual_engine
1688
+ previous_hidden_states = kwargs.get("previous_hidden_states")
1689
if prefill_meta is None and decode_meta.use_cuda_graph:
1690
assert model_input.input_tokens is not None
1691
graph_batch_size = model_input.input_tokens.shape[0]
1692
model_executable = self.graph_runners[virtual_engine][
1693
graph_batch_size]
1694
+ if previous_hidden_states is not None:
1695
+ previous_hidden_states = torch.cat([
1696
+ previous_hidden_states,
1697
+ torch.empty([
1698
+ graph_batch_size - previous_hidden_states.shape[0],
1699
+ *previous_hidden_states.shape[1:]
1700
+ ],
1701
+ dtype=previous_hidden_states.dtype,
1702
+ device=previous_hidden_states.device)
1703
+ ])
1704
else:
1705
model_executable = self.model
1706
@@ -1716,7 +1727,6 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
1727
"finished_requests_ids": model_input.finished_requests_ids,
1728
"request_ids_to_seq_ids": model_input.request_ids_to_seq_ids,
1729
} if self.has_inner_state else {}
1719
- previous_hidden_states = kwargs.get("previous_hidden_states")
1730
model_kwargs = {}
1731
if previous_hidden_states is not None:
1732
model_kwargs["previous_hidden_states"] = previous_hidden_states