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, 12:15 UTC
Fix more broken speculative decode tests (#17450)
Signed-off-by: Huy Do <huydhn@gmail.com>
Huy Do committed
May 1, 2025 at 06:05 UTC
b74d888c6357b94f3a5caf857a55af0c5886b731
4 files changed
+9
-4
tests/spec_decode/e2e/test_medusa_correctness.py
+1
-1
@@ -205,7 +205,7 @@ def test_medusa_e2e_greedy_correctness_cuda_graph(
205
@pytest.mark.parametrize(
206
"common_llm_kwargs",
207
[{
208
- "block_size": 8,
208
+ "block_size": 16,
209
# 2 for small prompt, 256//8 for generated.
210
"num_gpu_blocks_override": 2 + 256 // 8,
211
"max_model_len": (2 + 256 // 8) * 8,
tests/spec_decode/e2e/test_mlp_correctness.py
+2
-2
@@ -267,7 +267,7 @@ def test_mlp_e2e_seeded_correctness(vllm_runner, common_llm_kwargs,
267
@pytest.mark.parametrize(
268
"common_llm_kwargs",
269
[{
270
- "block_size": 8,
270
+ "block_size": 16,
271
# 2 for small prompt, 256//8 for generated.
272
"num_gpu_blocks_override": 2 + 256 // 8,
273
"max_model_len": (2 + 256 // 8) * 8,
@@ -321,7 +321,7 @@ def test_mlp_e2e_greedy_correctness_with_preemption(
321
@pytest.mark.parametrize(
322
"common_llm_kwargs",
323
[{
324
- "block_size": 8,
324
+ "block_size": 16,
325
# 2 for small prompt, 256//8 for generated.
326
"num_gpu_blocks_override": 2 + 256 // 8,
327
"max_model_len": (2 + 256 // 8) * 8,
tests/spec_decode/e2e/test_ngram_correctness.py
+1
-1
@@ -152,7 +152,7 @@ def test_ngram_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
152
@pytest.mark.parametrize(
153
"common_llm_kwargs",
154
[{
155
- "block_size": 8,
155
+ "block_size": 16,
156
# 2 for small prompt, 256//8 for generated.
157
"num_gpu_blocks_override": 2 + 256 // 8,
158
"max_model_len": (2 + 256 // 8) * 8,
vllm/spec_decode/multi_step_worker.py
+5
@@ -51,9 +51,14 @@ class MultiStepWorker(ProposerWorkerBase, DelegateWorkerBase):
51
def set_include_gpu_probs_tensor(self) -> None:
52
# Need include_gpu_probs_tensor for MultiStepWorker
53
self.model_runner.sampler.include_gpu_probs_tensor = True
54
+ if hasattr(self.model_runner.model, "sampler"):
55
+ (self.model_runner.model.sampler.include_gpu_probs_tensor) = True
56
57
def set_should_modify_greedy_probs_inplace(self) -> None:
58
self.model_runner.sampler.should_modify_greedy_probs_inplace = True
59
+ if hasattr(self.model_runner.model, "sampler"):
60
+ (self.model_runner.model.sampler.should_modify_greedy_probs_inplace
61
+ ) = True
62
63
@torch.inference_mode()
64
def sampler_output(