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, 09:45 UTC
[Spec Decode][V0] Fix spec decode correctness test in V0 eagle/medusa (#18175)
Signed-off-by: wwl2755 <wangwenlong2755@gmail.com>
wwl2755 committed
May 18, 2025 at 21:49 UTC
9da1095daf1710a41da3f79f7954f5dee956ce62
4 files changed
+21
-3
tests/spec_decode/e2e/test_eagle_correctness.py
-2
@@ -178,8 +178,6 @@ def test_eagle_e2e_greedy_correctness_cuda_graph(
178
batch_size, output_len, seed)
179
180
181
-# TRACKING: https://github.com/vllm-project/vllm/issues/18166
182
-@pytest.mark.skip(reason="RE-ENABLE: Failing on main.")
181
@pytest.mark.parametrize(
182
"common_llm_kwargs",
183
[{
vllm/model_executor/models/eagle.py
+11
@@ -146,6 +146,17 @@ class EAGLE(nn.Module):
146
if inputs_embeds is None:
147
inputs_embeds = self.get_input_embeddings(input_ids)
148
149
+ # Handle both empty previous_hidden_states
150
+ # and mismatched batch size
151
+ batch_size = inputs_embeds.size(0)
152
+ if previous_hidden_states.size(0) == 0 or \
153
+ previous_hidden_states.size(0) != batch_size:
154
+ hidden_dim = self.config.model.hidden_size
155
+ device = inputs_embeds.device
156
+ # Create zero tensor with matching batch size
157
+ previous_hidden_states = \
158
+ torch.zeros(batch_size, hidden_dim, device=device)
159
+
160
if self.add_para_norm:
161
inputs_embeds = torch.cat([
162
self.enorm(inputs_embeds),
vllm/model_executor/models/medusa.py
+8
-1
@@ -164,7 +164,14 @@ class Medusa(nn.Module):
164
self,
165
previous_hidden_states: torch.Tensor,
166
sampling_metadata: SamplingMetadata,
167
- ) -> list[SamplerOutput]:
167
+ ) -> Optional[list[SamplerOutput]]:
168
+ # During preemption, we may receive an empty tensor (batch_size=0)
169
+ if previous_hidden_states.size(0) == 0:
170
+ # Return None to signal the Top1Proposer that no proposals
171
+ # were generated for this batch, allowing it to handle this
172
+ # special case appropriately
173
+ return None
174
+
175
return self.sample(
176
logits=self.compute_logits(
177
hidden_states=self.forward(previous_hidden_states),
vllm/sequence.py
+2
@@ -1330,6 +1330,8 @@ class HiddenStates(msgspec.Struct, array_like=True,
1330
# may be "paused" then "resumed" later. This should only prune sequences
1331
# which are confirmed to be aborted.
1332
seq_ids = get_all_seq_ids(seq_group_metadata_list)
1333
+ # Only keep sequence IDs that exist in self._seq_ids
1334
+ seq_ids = [seq_id for seq_id in seq_ids if seq_id in self._seq_ids]
1335
if seq_ids != self._seq_ids:
1336
# Batch contents changed - prune removed sequences.
1337
index = [self._seq_ids.index(seq_id) for seq_id in seq_ids]