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, 01:00 UTC
[V0 Deprecation] Remove V0 Spec Decode workers (#21152)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jul 18, 2025 at 21:47 UTC
dd572c0ab3effa539b74f9a1288bb61ce83ada76
73 files changed
+191
-14275
.buildkite/test-pipeline.yaml
-14
@@ -159,7 +159,6 @@ steps:
159
- tests/distributed/test_utils
160
- tests/distributed/test_pynccl
161
- tests/distributed/test_events
162
- - tests/spec_decode/e2e/test_integration_dist_tp4
162
- tests/compile/test_basic_correctness
163
- examples/offline_inference/rlhf.py
164
- examples/offline_inference/rlhf_colocate.py
@@ -182,7 +181,6 @@ steps:
181
- pytest -v -s compile/test_basic_correctness.py
182
- pytest -v -s distributed/test_pynccl.py
183
- pytest -v -s distributed/test_events.py
185
- - pytest -v -s spec_decode/e2e/test_integration_dist_tp4.py
184
# TODO: create a dedicated test section for multi-GPU example tests
185
# when we have multiple distributed example tests
186
- pushd ../examples/offline_inference
@@ -330,17 +328,6 @@ steps:
328
- pytest -v -s samplers
329
- VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
330
333
-- label: Speculative decoding tests # 40min
334
- mirror_hardwares: [amdexperimental]
335
- source_file_dependencies:
336
- - vllm/spec_decode
337
- - tests/spec_decode
338
- - vllm/model_executor/models/eagle.py
339
- commands:
340
- - pytest -v -s spec_decode/e2e/test_multistep_correctness.py
341
- - VLLM_ATTENTION_BACKEND=FLASH_ATTN pytest -v -s spec_decode --ignore=spec_decode/e2e/test_multistep_correctness.py --ignore=spec_decode/e2e/test_mtp_correctness.py
342
- - pytest -v -s spec_decode/e2e/test_eagle_correctness.py
343
-
331
- label: LoRA Test %N # 15min each
332
mirror_hardwares: [amdexperimental, amdproduction]
333
source_file_dependencies:
@@ -726,7 +713,6 @@ steps:
713
- pytest -v -s distributed/test_sequence_parallel.py
714
# this test fails consistently.
715
# TODO: investigate and fix
729
- # - pytest -v -s spec_decode/e2e/test_integration_dist_tp2.py
716
- VLLM_USE_V1=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s test_sharded_state_loader.py
717
- VLLM_USE_V1=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s kv_transfer/test_disagg.py
718
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
.github/CODEOWNERS
-1
@@ -43,7 +43,6 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
43
/tests/multimodal @DarkLight1337 @ywang96
44
/tests/prefix_caching @comaniac @KuntaiDu
45
/tests/quantization @mgoin @robertgshaw2-redhat
46
-/tests/spec_decode @njhill @LiuXiaoxuanPKU
46
/tests/test_inputs.py @DarkLight1337 @ywang96
47
/tests/v1/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
48
/tests/v1/structured_output @mgoin @russellb @aarnphm
.github/mergify.yml
-3
@@ -164,10 +164,7 @@ pull_request_rules:
164
description: Automatically apply speculative-decoding label
165
conditions:
166
- or:
167
- - files~=^vllm/spec_decode/
167
- files~=^vllm/v1/spec_decode/
169
- - files=vllm/model_executor/layers/spec_decode_base_sampler.py
170
- - files~=^tests/spec_decode/
168
- files~=^tests/v1/spec_decode/
169
- files~=^examples/.*(spec_decode|mlpspeculator|eagle|speculation).*\.py
170
- files~=^vllm/model_executor/models/.*eagle.*\.py
pyproject.toml
-1
@@ -73,7 +73,6 @@ line-length = 80
73
"vllm/engine/**/*.py" = ["UP006", "UP035"]
74
"vllm/executor/**/*.py" = ["UP006", "UP035"]
75
"vllm/prompt_adapter/**/*.py" = ["UP006", "UP035"]
76
-"vllm/spec_decode/**/*.py" = ["UP006", "UP035"]
76
"vllm/worker/**/*.py" = ["UP006", "UP035"]
77
# Python 3.8 typing - skip utils for ROCm
78
"vllm/utils/__init__.py" = ["UP006", "UP035"]
tests/core/test_serialization.py
+1
-1
@@ -6,7 +6,7 @@ import msgspec
6
from vllm.executor.msgspec_utils import decode_hook, encode_hook
7
from vllm.sequence import ExecuteModelRequest
8
9
-from ..spec_decode.utils import create_batch
9
+from .utils import create_batch
10
11
12
def test_msgspec_serialization():
tests/core/utils.py
+131
-3
@@ -4,15 +4,16 @@
4
import time
5
from collections import defaultdict
6
from collections.abc import Sequence as GenericSequence
7
-from typing import Any, Optional
7
+from itertools import count
8
+from typing import Any, Optional, Union
9
10
import torch
11
11
-from vllm import SamplingParams
12
from vllm.core.scheduler import Scheduler, SchedulerOutputs
13
from vllm.inputs import EncoderDecoderInputs, embeds_inputs, token_inputs
14
from vllm.lora.request import LoRARequest
15
-from vllm.sequence import (Logprob, Sequence, SequenceGroup,
15
+from vllm.sampling_params import SamplingParams
16
+from vllm.sequence import (Logprob, Sequence, SequenceData, SequenceGroup,
17
SequenceGroupMetadata)
18
19
@@ -262,3 +263,130 @@ class SchedulerProxy:
263
self, ) -> tuple[list[SequenceGroupMetadata], SchedulerOutputs, Any]:
264
_, _, ret = self.call_history["schedule"][-1]
265
return ret
266
+
267
+
268
+def create_seq_group_metadata_from_prompts(
269
+ prompts: list[list[int]],
270
+ num_gpu_blocks: int,
271
+ block_size: int,
272
+ final_prompt_lens: list[int],
273
+ continuations: Optional[list[list[int]]] = None,
274
+ seq_ids: Optional[list[int]] = None,
275
+) -> list[SequenceGroupMetadata]:
276
+
277
+ if continuations is None:
278
+ continuations = [[] for _ in prompts]
279
+
280
+ if seq_ids is None:
281
+ seq_ids = list(i for i, _ in enumerate(prompts))
282
+
283
+ free_gpu_blocks = list(range(num_gpu_blocks))
284
+
285
+ block_allocations = {
286
+ i: [
287
+ free_gpu_blocks.pop()
288
+ for _ in range(round_up_to_next_block(final_len, block_size))
289
+ ]
290
+ for i, final_len in enumerate(final_prompt_lens)
291
+ }
292
+
293
+ seq_grou_metadata_list = []
294
+ for i, (prompt_token_ids,
295
+ cont_token_ids) in enumerate(zip(prompts, continuations)):
296
+ data = SequenceData.from_seqs(prompt_token_ids, cont_token_ids)
297
+ data.update_num_computed_tokens(
298
+ len(prompt_token_ids) + len(cont_token_ids) - 1)
299
+ seq_data = {i: data}
300
+ seq_grou_metadata_list.append(
301
+ SequenceGroupMetadata(
302
+ request_id=str(i),
303
+ is_prompt=len(cont_token_ids) == 0,
304
+ seq_data=seq_data,
305
+ sampling_params=SamplingParams(temperature=0.0),
306
+ block_tables={i: block_allocations[i][:]},
307
+ ))
308
+ return seq_grou_metadata_list
309
+
310
+
311
+def create_chunked_seq_group_metadata_from_prompt(
312
+ prompt: list[int],
313
+ num_gpu_blocks: int,
314
+ chunk_size: int,
315
+ block_size: int,
316
+ seq_id: Optional[int] = None) -> list[SequenceGroupMetadata]:
317
+
318
+ if seq_id is None:
319
+ seq_id = 0
320
+
321
+ free_gpu_blocks = list(range(num_gpu_blocks))
322
+
323
+ block_allocations = [
324
+ free_gpu_blocks.pop()
325
+ for _ in range(round_up_to_next_block(len(prompt), block_size))
326
+ ]
327
+
328
+ seq_group_metadata_list = []
329
+ for i, idx in enumerate(range(0, len(prompt), chunk_size)):
330
+ chunk_ids = prompt[idx:idx + chunk_size]
331
+ data = SequenceData.from_seqs(prompt)
332
+ data.update_num_computed_tokens(idx)
333
+ seq_data = {i: data}
334
+ seq_group_metadata_list.append(
335
+ SequenceGroupMetadata(
336
+ request_id=str(seq_id),
337
+ is_prompt=True,
338
+ do_sample=idx + chunk_size >= len(prompt), # terminal chunk
339
+ seq_data=seq_data,
340
+ sampling_params=SamplingParams(temperature=0.0),
341
+ block_tables={i: block_allocations},
342
+ token_chunk_size=len(chunk_ids)))
343
+ return seq_group_metadata_list
344
+
345
+
346
+def create_batch(batch_size,
347
+ k,
348
+ prompt_len: Union[int, list[int]] = 10,
349
+ prev_output_token_len: int = 10,
350
+ seq_ids: Optional[list[int]] = None,
351
+ num_gpu_blocks: Optional[int] = None,
352
+ block_size: Optional[int] = None,
353
+ prefill_chunk_size: Optional[int] = None):
354
+ if block_size is None:
355
+ block_size = 8
356
+
357
+ if num_gpu_blocks is None:
358
+ num_gpu_blocks = 2048 // block_size
359
+
360
+ iterator = count()
361
+
362
+ if isinstance(prompt_len, int):
363
+ prompt_lens = [prompt_len for _ in range(batch_size)]
364
+ else:
365
+ prompt_lens = prompt_len
366
+
367
+ prompts = [[next(iterator) for _ in range(p_len)] for p_len in prompt_lens]
368
+
369
+ if prefill_chunk_size:
370
+ # Create a batch of chunked prompts.
371
+ if not seq_ids:
372
+ seq_ids = list(range(len(prompts)))
373
+ seq_group_metadata_list = []
374
+ for p, sid in zip(prompts, seq_ids):
375
+ seq_group_metadata_list += \
376
+ create_chunked_seq_group_metadata_from_prompt(
377
+ p, num_gpu_blocks, prefill_chunk_size, block_size, sid)
378
+ seq_group_metadata_list = seq_group_metadata_list[:batch_size]
379
+ prev_output_tokens = []
380
+ else:
381
+ prev_output_tokens = [[
382
+ next(iterator) for _ in range(prev_output_token_len)
383
+ ] for _ in range(batch_size)]
384
+ final_prompt_lens = [
385
+ len(prompt) + len(prev_output_token) + k + 1
386
+ for prompt, prev_output_token in zip(prompts, prev_output_tokens)
387
+ ]
388
+
389
+ seq_group_metadata_list = create_seq_group_metadata_from_prompts(
390
+ prompts, num_gpu_blocks, block_size, final_prompt_lens,
391
+ prev_output_tokens, seq_ids)
392
+ return seq_group_metadata_list, prompts, prev_output_tokens
tests/metrics/test_metrics.py
-146
@@ -1,15 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import time
5
-
4
import pytest
5
import ray
6
from prometheus_client import REGISTRY
7
8
import vllm.envs as envs
9
from vllm import EngineArgs, LLMEngine
12
-from vllm.distributed import cleanup_dist_env_and_memory
10
from vllm.engine.arg_utils import AsyncEngineArgs
11
from vllm.engine.async_llm_engine import AsyncLLMEngine
12
from vllm.engine.metrics import RayPrometheusStatLogger
@@ -232,149 +229,6 @@ def test_engine_log_metrics_regression(
229
assert_metrics(model, engine, disable_log_stats, len(example_prompts))
230
231
235
-@pytest.mark.parametrize("model", MODELS)
236
-@pytest.mark.parametrize("dtype", ["half"])
237
-@pytest.mark.parametrize("max_tokens", [10])
238
-def test_metric_spec_decode(
239
- vllm_runner,
240
- example_prompts,
241
- model: str,
242
- dtype: str,
243
- max_tokens: int,
244
-) -> None:
245
- k = 5
246
-
247
- with vllm_runner(
248
- model,
249
- dtype=dtype,
250
- disable_log_stats=False,
251
- gpu_memory_utilization=0.4,
252
- speculative_config={
253
- "model": model,
254
- "num_speculative_tokens": k,
255
- },
256
- ) as vllm_model:
257
-
258
- # Force log interval to be 0 to catch all metrics.
259
- stat_logger = vllm_model.model.llm_engine.stat_loggers['prometheus']
260
- stat_logger.local_interval = 0
261
-
262
- # Note that the purpose of this test is to verify spec decode
263
- # metrics instead of functional correctness, so the expected values
264
- # are intended to be loose.
265
- metric_name_to_expected_fn = {
266
- "gauge_spec_decode_draft_acceptance_rate": lambda v: 0 <= v <= 1,
267
- "gauge_spec_decode_efficiency": lambda v: 0 <= v <= 1,
268
- "counter_spec_decode_num_accepted_tokens": lambda v: 0 <= v <= k,
269
- "counter_spec_decode_num_draft_tokens": lambda v: v == k,
270
- "counter_spec_decode_num_emitted_tokens":
271
- lambda v: 0 <= v <= k + 1,
272
- }
273
-
274
- # Use one request to better inspect the metrics.
275
- prompts = example_prompts[:1]
276
-
277
- _ = vllm_model.generate_greedy(prompts, max_tokens)
278
- for metric_name, is_expected in metric_name_to_expected_fn.items():
279
- metric_val = getattr(
280
- stat_logger.metrics,
281
- metric_name).labels(**stat_logger.labels)._value.get()
282
- assert is_expected(metric_val), (
283
- f"the value of metric {metric_name} ({metric_val}) "
284
- "does not meet expectation")
285
-
286
-
287
-@pytest.mark.parametrize("model", MODELS)
288
-@pytest.mark.parametrize("dtype", ["half"])
289
-@pytest.mark.parametrize("max_tokens", [10])
290
-@pytest.mark.parametrize("log_interval", [1, 3, 5, 7])
291
-def test_metric_spec_decode_interval(
292
- vllm_runner,
293
- example_prompts,
294
- model: str,
295
- dtype: str,
296
- max_tokens: int,
297
- log_interval: int,
298
-) -> None:
299
- k = 5
300
-
301
- engine_args = EngineArgs(
302
- model=model,
303
- dtype=dtype,
304
- disable_log_stats=False,
305
- gpu_memory_utilization=0.4,
306
- speculative_config={
307
- "model": model,
308
- "num_speculative_tokens": k,
309
- },
310
- enforce_eager=True,
311
- )
312
-
313
- engine = LLMEngine.from_engine_args(engine_args)
314
-
315
- try:
316
-
317
- engine.add_request(
318
- "request-id-0",
319
- example_prompts[0],
320
- SamplingParams(max_tokens=max_tokens),
321
- )
322
-
323
- # set log internal
324
- stat_logger = engine.stat_loggers['prometheus']
325
- stat_logger.local_interval = log_interval
326
-
327
- # prefill
328
- engine.step()
329
-
330
- # wait for 5 seconds to ensure that spec decode metrics
331
- # get triggered in first decode step
332
- time.sleep(5)
333
-
334
- # first decode step should trigger async collection of metrics
335
- engine.step()
336
-
337
- # wait one second to allow H2D transfer to finish
338
- time.sleep(1)
339
-
340
- # second decode step should now be able to collect the spec
341
- # decode stats and the request should also be finished
342
- engine.step()
343
-
344
- # must have finisehd now
345
- assert not engine.has_unfinished_requests()
346
-
347
- # wait to ensure logging occurs
348
- time.sleep(log_interval)
349
-
350
- # force logging
351
- engine.step()
352
-
353
- # Note that the purpose of this test is to verify spec decode
354
- # metrics instead of functional correctness, so the expected values
355
- # are intended to be loose.
356
- metric_name_to_expected_fn = {
357
- "gauge_spec_decode_draft_acceptance_rate": lambda v: 0 <= v <= 1,
358
- "gauge_spec_decode_efficiency": lambda v: 0 <= v <= 1,
359
- "counter_spec_decode_num_accepted_tokens": lambda v: 0 <= v <= k,
360
- "counter_spec_decode_num_draft_tokens": lambda v: v == k,
361
- "counter_spec_decode_num_emitted_tokens":
362
- lambda v: 0 <= v <= k + 1,
363
- }
364
-
365
- for metric_name, is_expected in metric_name_to_expected_fn.items():
366
- metric_val = getattr(
367
- stat_logger.metrics,
368
- metric_name).labels(**stat_logger.labels)._value.get()
369
- assert is_expected(metric_val), (
370
- f"the value of metric {metric_name} ({metric_val}) "
371
- "does not meet expectation")
372
-
373
- finally:
374
- del engine
375
- cleanup_dist_env_and_memory()
376
-
377
-
232
def assert_metrics(model: str, engine: LLMEngine, disable_log_stats: bool,
233
num_requests: int) -> None:
234
if disable_log_stats:
tests/models/registry.py
+4
-4
@@ -457,12 +457,12 @@ _MULTIMODAL_EXAMPLE_MODELS = {
457
458
459
_SPECULATIVE_DECODING_EXAMPLE_MODELS = {
460
- "EAGLEModel": _HfExamplesInfo("JackFram/llama-68m",
461
- speculative_model="abhigoyal/vllm-eagle-llama-68m-random"), # noqa: E501
460
"MedusaModel": _HfExamplesInfo("JackFram/llama-68m",
461
speculative_model="abhigoyal/vllm-medusa-llama-68m-random"), # noqa: E501
464
- "MLPSpeculatorPreTrainedModel": _HfExamplesInfo("JackFram/llama-160m",
465
- speculative_model="ibm-ai-platform/llama-160m-accelerator"), # noqa: E501
462
+ # Temporarily disabled.
463
+ # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.
464
+ # "MLPSpeculatorPreTrainedModel": _HfExamplesInfo("JackFram/llama-160m",
465
+ # speculative_model="ibm-ai-platform/llama-160m-accelerator"), # noqa: E501
466
"DeepSeekMTPModel": _HfExamplesInfo("luccafong/deepseek_mtp_main_random",
467
speculative_model="luccafong/deepseek_mtp_draft_random", # noqa: E501
468
trust_remote_code=True),
tests/models/test_registry.py
+9
-5
@@ -72,11 +72,15 @@ def test_registry_model_property(model_arch, is_mm, init_cuda, is_ce):
72
73
74
@create_new_process_for_each_test()
75
-@pytest.mark.parametrize("model_arch,is_pp,init_cuda", [
76
- ("MLPSpeculatorPreTrainedModel", False, False),
77
- ("DeepseekV2ForCausalLM", True, False),
78
- ("Qwen2VLForConditionalGeneration", True, True),
79
-])
75
+@pytest.mark.parametrize(
76
+ "model_arch,is_pp,init_cuda",
77
+ [
78
+ # TODO(woosuk): Re-enable this once the MLP Speculator is supported
79
+ # in V1.
80
+ # ("MLPSpeculatorPreTrainedModel", False, False),
81
+ ("DeepseekV2ForCausalLM", True, False),
82
+ ("Qwen2VLForConditionalGeneration", True, True),
83
+ ])
84
def test_registry_is_pp(model_arch, is_pp, init_cuda):
85
assert ModelRegistry.is_pp_supported_model(model_arch) is is_pp
86
tests/samplers/test_rejection_sampler.py
deleted
-577
@@ -1,577 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests for rejection sampling."""
4
-
5
-import pytest
6
-import torch
7
-import torch.nn.functional as F
8
-
9
-from vllm.model_executor.layers.rejection_sampler import RejectionSampler
10
-from vllm.model_executor.utils import set_random_seed
11
-
12
-
13
-@pytest.fixture(scope="function", autouse=True)
14
-def use_v0_only(monkeypatch):
15
- """
16
- This file tests V0 internals, so set VLLM_USE_V1=0.
17
- """
18
- monkeypatch.setenv('VLLM_USE_V1', '0')
19
-
20
-
21
-CUDA_DEVICES = [
22
- f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)
23
-]
24
-
25
-
26
-def mock_causal_accepted_tensor(
27
- k: int, last_accepted_indices: torch.Tensor) -> torch.Tensor:
28
- """Generate an "accepted" tensor which should yield causally-accepted tokens
29
- up to last accepted indices.
30
-
31
- Tokens after last_accepted_indices+1 may also be accepted, although they
32
- will not be causally accepted.
33
- """
34
- batch_size = last_accepted_indices.shape[0]
35
-
36
- accepted = (torch.arange(k).expand(batch_size, k)
37
- <= last_accepted_indices.unsqueeze(-1).broadcast_to(
38
- batch_size, k))
39
-
40
- # Sprinkle accepted values after the contiguous initial accepted values.
41
- # This replicates the behavior of rejection sampling, which may "accept"
42
- # a token that cannot be accepted because of causality.
43
- sprinkle_candidates = (torch.arange(k).expand(
44
- batch_size,
45
- k) > last_accepted_indices.unsqueeze(-1).broadcast_to(batch_size, k) +
46
- 1)
47
- sprinkle = torch.rand(batch_size, k) > 0.5
48
- accepted[sprinkle_candidates] = sprinkle[sprinkle_candidates]
49
- return accepted
50
-
51
-
52
-@pytest.mark.parametrize("seed", list(range(10)))
53
-@pytest.mark.parametrize(
54
- "which_tokens_accepted",
55
- ["all_tokens_accepted", "no_tokens_accepted", "some_tokens_accepted"])
56
-@pytest.mark.parametrize("device", CUDA_DEVICES)
57
-@pytest.mark.parametrize("use_flashinfer", [True, False])
58
-@torch.inference_mode()
59
-def test_correct_output_format(which_tokens_accepted: str, seed: int,
60
- device: str, use_flashinfer: bool):
61
- """Verify the output has correct format given predetermined accepted matrix.
62
- """
63
- set_random_seed(seed)
64
- torch.set_default_device(device)
65
-
66
- batch_size = 10
67
- k = 5
68
- vocab_size = 3000
69
-
70
- if which_tokens_accepted == "all_tokens_accepted":
71
- accepted = mock_causal_accepted_tensor(
72
- k, -1 + k * torch.ones((batch_size, ), dtype=torch.long))
73
- elif which_tokens_accepted == "no_tokens_accepted":
74
- accepted = mock_causal_accepted_tensor(
75
- k, -torch.ones((batch_size, ), dtype=torch.long))
76
- elif which_tokens_accepted == "some_tokens_accepted":
77
- last_accepted_indices = torch.randint(low=-1,
78
- high=k,
79
- size=(batch_size, ))
80
- accepted = mock_causal_accepted_tensor(k, last_accepted_indices)
81
- else:
82
- raise AssertionError()
83
-
84
- recovered_token_ids = torch.randint(low=0,
85
- high=vocab_size,
86
- size=(batch_size, k),
87
- dtype=torch.int64)
88
- draft_token_ids = torch.randint(low=0,
89
- high=vocab_size,
90
- size=(batch_size, k),
91
- dtype=torch.int64)
92
- bonus_token_ids = torch.randint(low=0,
93
- high=vocab_size,
94
- size=(batch_size, 1),
95
- dtype=torch.int64)
96
-
97
- rejection_sampler = RejectionSampler(use_flashinfer=use_flashinfer)
98
- rejection_sampler.init_gpu_tensors(device=device)
99
- output_token_ids = rejection_sampler._create_output( # pylint: disable=protected-access
100
- accepted,
101
- recovered_token_ids,
102
- draft_token_ids,
103
- bonus_token_ids,
104
- )
105
-
106
- expected_bonus_token_ids = bonus_token_ids.clone()
107
-
108
- if which_tokens_accepted == "all_tokens_accepted":
109
- # Expect all tokens to be equal to draft tokens.
110
- assert torch.equal(output_token_ids[:, :-1], draft_token_ids)
111
-
112
- # Expect all bonus tokens to be included.
113
- assert torch.equal(output_token_ids[:, -1:], expected_bonus_token_ids)
114
- elif which_tokens_accepted == "no_tokens_accepted":
115
- # Expect first token to be equal to recovered tokens.
116
- assert torch.equal(output_token_ids[:, 0], recovered_token_ids[:, 0])
117
-
118
- # Expect everything else to be -1.
119
- assert torch.equal(output_token_ids[:, 1:],
120
- torch.ones_like(output_token_ids[:, 1:]) * -1)
121
- elif which_tokens_accepted == "some_tokens_accepted":
122
- recovered_plus_bonus = torch.cat(
123
- (recovered_token_ids, expected_bonus_token_ids), dim=-1)
124
- # Assert first rejected token is a recovered token or bonus token.
125
- assert torch.equal(
126
- recovered_plus_bonus[torch.arange(0, batch_size),
127
- last_accepted_indices + 1],
128
- output_token_ids[torch.arange(0, batch_size),
129
- last_accepted_indices + 1])
130
-
131
- # Assert every subsequent token is -1.
132
- subsequent_mask = torch.arange(0, k + 1).expand(
133
- batch_size, k + 1) >= (last_accepted_indices + 2).unsqueeze(-1)
134
- assert torch.all(output_token_ids[subsequent_mask] == -1)
135
-
136
-
137
-@pytest.mark.parametrize("k", list(range(1, 6)))
138
-@pytest.mark.parametrize("vocab_size", [30_000, 50_000])
139
-@pytest.mark.parametrize("batch_size", list(range(1, 32)))
140
-@pytest.mark.parametrize("device", CUDA_DEVICES)
141
-@pytest.mark.parametrize("use_flashinfer", [True, False])
142
-@torch.inference_mode()
143
-def test_no_crash_with_varying_dims(k: int, vocab_size: int, batch_size: int,
144
- device: str, use_flashinfer: bool):
145
- torch.set_default_device(device)
146
- rejection_sampler = RejectionSampler(use_flashinfer=use_flashinfer)
147
- rejection_sampler.init_gpu_tensors(device=device)
148
-
149
- draft_probs = torch.rand(batch_size, k, vocab_size, dtype=torch.float32)
150
- target_probs = torch.rand(batch_size,
151
- k + 1,
152
- vocab_size,
153
- dtype=torch.float32)
154
- bonus_token_ids = torch.randint(low=0,
155
- high=vocab_size,
156
- size=(batch_size, 1),
157
- dtype=torch.int64)
158
- draft_token_ids = torch.randint(low=0,
159
- high=vocab_size,
160
- size=(batch_size, k),
161
- dtype=torch.int64)
162
-
163
- rejection_sampler(target_probs, bonus_token_ids, draft_probs,
164
- draft_token_ids)
165
-
166
-
167
-@pytest.mark.parametrize("frac_seeded", [0.0, 0.25, 0.5, 1.0])
168
-@pytest.mark.parametrize("k", [1, 3, 6])
169
-@pytest.mark.parametrize("vocab_size", [30_000, 50_000])
170
-@pytest.mark.parametrize("batch_size", [1, 8, 32, 128])
171
-@pytest.mark.parametrize("n_rep", [100])
172
-@pytest.mark.parametrize("device", CUDA_DEVICES)
173
-# @pytest.mark.parametrize("use_flashinfer", [True, False])
174
-# Not testing FlashInfer now, since 0.2.3 API removed the ability
175
-# to pass in uniform samples.
176
-@pytest.mark.parametrize("use_flashinfer", [False])
177
-@torch.inference_mode()
178
-def test_deterministic_when_seeded(k: int, vocab_size: int, batch_size: int,
179
- frac_seeded: float, n_rep: int, device: str,
180
- use_flashinfer: bool):
181
- torch.set_default_device(device)
182
- rejection_sampler = RejectionSampler(use_flashinfer=use_flashinfer)
183
- rejection_sampler.init_gpu_tensors(device=device)
184
-
185
- draft_probs = torch.rand(batch_size, k, vocab_size, dtype=torch.float32)
186
- target_probs = torch.rand(batch_size,
187
- k + 1,
188
- vocab_size,
189
- dtype=torch.float32)
190
- bonus_token_ids = torch.randint(low=0,
191
- high=vocab_size,
192
- size=(batch_size, 1),
193
- dtype=torch.int64)
194
- draft_token_ids = torch.randint(low=0,
195
- high=vocab_size,
196
- size=(batch_size, k),
197
- dtype=torch.int64)
198
-
199
- seeded_mask = torch.rand(batch_size, dtype=torch.float32) <= frac_seeded
200
-
201
- results = []
202
- for _ in range(n_rep):
203
- seeded_seqs = {
204
- i: torch.Generator(device=device).manual_seed(i)
205
- for i in range(batch_size) if seeded_mask[i]
206
- }
207
- results.append(
208
- rejection_sampler(target_probs, bonus_token_ids, draft_probs,
209
- draft_token_ids, seeded_seqs))
210
-
211
- for i in range(batch_size):
212
- if seeded_mask[i]:
213
- for j in range(1, n_rep):
214
- assert torch.equal(results[j][i], results[0][i])
215
-
216
-
217
-@pytest.mark.parametrize("k", [1, 3, 6])
218
-@pytest.mark.parametrize("vocab_size", [30_000, 50_000])
219
-@pytest.mark.parametrize("batch_size", [3, 8, 32, 128])
220
-@pytest.mark.parametrize("device", CUDA_DEVICES)
221
-# @pytest.mark.parametrize("use_flashinfer", [True, False])
222
-# Not testing FlashInfer now, since 0.2.3 API removed the ability
223
-# to pass in uniform samples.
224
-@pytest.mark.parametrize("use_flashinfer", [False])
225
-@torch.inference_mode()
226
-def test_mixed_seeded_batch(k: int, vocab_size: int, batch_size: int,
227
- device: str, use_flashinfer: bool):
228
- torch.set_default_device(device)
229
- set_random_seed(0)
230
- draft_probs = torch.rand(batch_size, k, vocab_size, dtype=torch.float32)
231
- target_probs = torch.rand(batch_size,
232
- k + 1,
233
- vocab_size,
234
- dtype=torch.float32)
235
- bonus_token_ids = torch.randint(low=0,
236
- high=vocab_size,
237
- size=(batch_size, 1),
238
- dtype=torch.int64)
239
- draft_token_ids = torch.randint(low=0,
240
- high=vocab_size,
241
- size=(batch_size, k),
242
- dtype=torch.int64)
243
-
244
- single_batches = []
245
- for i in range(batch_size):
246
- single_batches.append((draft_probs[i].clone().unsqueeze(0),
247
- draft_token_ids[i].clone().unsqueeze(0),
248
- target_probs[i].clone().unsqueeze(0),
249
- bonus_token_ids[i].clone().unsqueeze(0),
250
- draft_token_ids[i].clone().unsqueeze(0)))
251
-
252
- set_random_seed(0)
253
- rejection_sampler = RejectionSampler(use_flashinfer=use_flashinfer)
254
- rejection_sampler.init_gpu_tensors(device=device)
255
-
256
- results = []
257
- seeded_seqs = {
258
- i: torch.Generator(device=device).manual_seed(i)
259
- for i in range(1, batch_size) # 0 is seed None
260
- }
261
- batch_result = rejection_sampler(target_probs.clone(),
262
- bonus_token_ids.clone(),
263
- draft_probs.clone(),
264
- draft_token_ids.clone(), seeded_seqs)
265
-
266
- set_random_seed(0)
267
-
268
- rejection_sampler = RejectionSampler(use_flashinfer=use_flashinfer)
269
- rejection_sampler.init_gpu_tensors(device=device)
270
- for i in range(batch_size):
271
- request_seeded_seqs = {
272
- 0: torch.Generator(device=device).manual_seed(i)
273
- } if seeded_seqs.get(i) is not None else None
274
- (draft_probs, draft_token_ids, target_probs, bonus_token_ids,
275
- draft_token_ids) = single_batches[i]
276
- results.append(
277
- rejection_sampler(target_probs, bonus_token_ids, draft_probs,
278
- draft_token_ids, request_seeded_seqs))
279
- for i in range(batch_size):
280
- assert torch.equal(batch_result[i], results[i].squeeze(0))
281
-
282
-
283
-@pytest.mark.parametrize("k", [1, 3, 6])
284
-@pytest.mark.parametrize("vocab_size", [30_000, 50_000])
285
-@pytest.mark.parametrize("batch_size", [1, 8, 32, 128])
286
-@pytest.mark.parametrize("device", CUDA_DEVICES)
287
-@torch.inference_mode()
288
-def test_compare_nonflashinfer_backend(k: int, vocab_size: int,
289
- batch_size: int, device: str):
290
- """
291
- Test the flashinfer and nonflashinfer backend generate
292
- the same output metrics.
293
- """
294
-
295
- pytest.skip("Not testing FlashInfer now, since 0.2.3 API removed "
296
- "the ability to pass in uniform samples.")
297
-
298
- torch.set_default_device(device)
299
- torch.manual_seed(0)
300
- draft_probs = torch.rand(batch_size, k, vocab_size, dtype=torch.float32)
301
- target_probs = torch.rand(batch_size,
302
- k + 1,
303
- vocab_size,
304
- dtype=torch.float32)
305
- bonus_token_ids = torch.randint(low=0,
306
- high=vocab_size,
307
- size=(batch_size, 1),
308
- dtype=torch.int64)
309
- draft_token_ids = torch.randint(low=0,
310
- high=vocab_size,
311
- size=(batch_size, k),
312
- dtype=torch.int64)
313
-
314
- num_accepted_tokens = []
315
- num_emitted_tokens = []
316
- num_draft_tokens = []
317
-
318
- def get_seeded_seqs():
319
- return {
320
- i: torch.Generator(device=device).manual_seed(i)
321
- for i in range(batch_size)
322
- }
323
-
324
- for use_flashinfer in [True, False]:
325
- rejection_sampler = RejectionSampler(use_flashinfer=use_flashinfer)
326
- rejection_sampler.init_gpu_tensors(device=device)
327
- # We use seeded sequences to ensure the same tokens are accepted
328
- # for both flashinfer and nonflashinfer backends.
329
- seeded_seqs = get_seeded_seqs()
330
- rejection_sampler(target_probs, bonus_token_ids, draft_probs,
331
- draft_token_ids, seeded_seqs)
332
- num_accepted_tokens.append(rejection_sampler.num_accepted_tokens)
333
- num_emitted_tokens.append(rejection_sampler.num_emitted_tokens)
334
- num_draft_tokens.append(rejection_sampler.num_draft_tokens)
335
-
336
- assert num_accepted_tokens[0] == num_accepted_tokens[1]
337
- assert num_emitted_tokens[0] == num_emitted_tokens[1]
338
- assert num_draft_tokens[0] == num_draft_tokens[1]
339
-
340
-
341
-@pytest.mark.parametrize("above_or_below_vocab_range", ["above", "below"])
342
-@pytest.mark.parametrize("which_token_ids",
343
- ["bonus_token_ids", "draft_token_ids"])
344
-@pytest.mark.parametrize("device", CUDA_DEVICES)
345
-@pytest.mark.parametrize("use_flashinfer", [True, False])
346
-@torch.inference_mode()
347
-def test_raises_when_vocab_oob(above_or_below_vocab_range: str,
348
- which_token_ids: str, device: str,
349
- use_flashinfer: bool):
350
- k = 3
351
- batch_size = 5
352
- vocab_size = 30_000
353
- torch.set_default_device(device)
354
-
355
- rejection_sampler = RejectionSampler(use_flashinfer=use_flashinfer,
356
- strict_mode=True)
357
- rejection_sampler.init_gpu_tensors(device=device)
358
-
359
- draft_probs = torch.rand(batch_size, k, vocab_size, dtype=torch.float32)
360
- target_probs = torch.rand(batch_size,
361
- k + 1,
362
- vocab_size,
363
- dtype=torch.float32)
364
- bonus_token_ids = torch.randint(low=0,
365
- high=vocab_size,
366
- size=(batch_size, 1),
367
- dtype=torch.int64)
368
- draft_token_ids = torch.randint(low=0,
369
- high=vocab_size,
370
- size=(batch_size, k),
371
- dtype=torch.int64)
372
-
373
- oob_token_ids = None
374
- if which_token_ids == "bonus_token_ids":
375
- oob_token_ids = bonus_token_ids
376
- elif which_token_ids == "draft_token_ids":
377
- oob_token_ids = draft_token_ids
378
- else:
379
- raise AssertionError()
380
-
381
- if above_or_below_vocab_range == "above":
382
- rogue_token_id = vocab_size + 1
383
- elif above_or_below_vocab_range == "below":
384
- rogue_token_id = -1
385
- else:
386
- raise AssertionError()
387
-
388
- oob_token_ids[0][0] = rogue_token_id
389
-
390
- with pytest.raises(AssertionError):
391
- rejection_sampler(target_probs, bonus_token_ids, draft_probs,
392
- draft_token_ids)
393
-
394
-
395
-@pytest.mark.parametrize("draft_and_target_probs_equal", [True, False])
396
-@pytest.mark.parametrize("seed", list(range(5)))
397
-@pytest.mark.parametrize("use_flashinfer", [True, False])
398
-@torch.inference_mode()
399
-def test_rejection_sampling_approximates_target_distribution(
400
- seed: int, draft_and_target_probs_equal: bool, use_flashinfer: bool):
401
- """Verify rejection sampling approximates target distribution,
402
- despite sampling from a potentially distinct draft distribution.
403
-
404
- This is done by first creating a random target probability
405
- distribution and a random draft probability distribution. We then
406
- sample token ids from the rejection sampler using these draft
407
- and target distributions. The samples are used to estimate
408
- the output probability distribution, which we expect to approximate
409
- the target distribution.
410
-
411
- A basic distance metric is used to determine similarity between
412
- distributions.
413
-
414
- We expect that as we increase the number of samples,
415
- the distance between the observed distribution and the target
416
- distribution decreases. To measure this, we compare the distance
417
- of the observed distribution against both the target distribution
418
- and a uniform random distribution. We expect the distance between
419
- the observed distribution and the target distribution to improve
420
- much more than the distance improvement between the observed
421
- distribution and the random distribution.
422
-
423
- When draft_and_target_probs_equal=True, the draft and target
424
- probabilities are exactly equal. Rejection sampling should
425
- still work without any NaNs or exceptions.
426
- """
427
- torch.set_default_device("cpu")
428
- set_random_seed(seed)
429
- helper = _CorrectnessTestHelper(
430
- vocab_size=10,
431
- rejection_sampler=RejectionSampler(use_flashinfer=use_flashinfer),
432
- )
433
-
434
- draft_probs, target_probs, reference_probs = helper.generate_probs_for_test(
435
- draft_and_target_probs_equal)
436
-
437
- sample_sizes = [10, 100, 1_000, 10_000, 100_000]
438
- distance_wrt_reference: list[float] = []
439
- distance_wrt_target: list[float] = []
440
-
441
- for num_samples in sample_sizes:
442
- (reference_vs_rejsample_dist,
443
- target_vs_rejsample_dist) = helper.run_and_compare_distributions(
444
- draft_probs,
445
- target_probs,
446
- reference_probs,
447
- num_samples,
448
- )
449
-
450
- distance_wrt_reference.append(reference_vs_rejsample_dist)
451
- distance_wrt_target.append(target_vs_rejsample_dist)
452
-
453
- relative_change_in_distance_wrt_target = get_ratio_first_to_last(
454
- distance_wrt_target)
455
- relative_change_in_distance_wrt_reference = get_ratio_first_to_last(
456
- distance_wrt_reference)
457
-
458
- print(f"{num_samples=} {target_vs_rejsample_dist=:.05f} "
459
- f"{reference_vs_rejsample_dist=:.05f}")
460
- print(f"{num_samples=} {relative_change_in_distance_wrt_target=:.02f} "
461
- f"{relative_change_in_distance_wrt_reference=:.02f}")
462
-
463
- relative_change_in_distance_wrt_target = get_ratio_first_to_last(
464
- distance_wrt_target)
465
- relative_change_in_distance_wrt_reference = get_ratio_first_to_last(
466
- distance_wrt_reference)
467
-
468
- expected_improvement_multiplier = 20
469
- assert (relative_change_in_distance_wrt_target
470
- > relative_change_in_distance_wrt_reference *
471
- expected_improvement_multiplier)
472
-
473
-
474
-def get_ratio_first_to_last(elements: list[float]) -> float:
475
- return elements[0] / elements[-1]
476
-
477
-
478
-class _CorrectnessTestHelper:
479
- """Class that packages together logic required for the unit-level
480
- rejection sampling correctness test.
481
- """
482
-
483
- def __init__(self, vocab_size: int, rejection_sampler: RejectionSampler):
484
- self.rejection_sampler = rejection_sampler
485
- self.vocab_size = vocab_size
486
- self.vocab_range = (0, vocab_size)
487
-
488
- self.rejection_sampler.init_gpu_tensors(device=0)
489
-
490
- # Keep test simple, use k=1
491
- self.k = 1
492
-
493
- # Bonus tokens not used, but rejection sampler requires
494
- # correct shape.
495
- self.num_bonus_tokens = 1
496
-
497
- def generate_probs_for_test(
498
- self, draft_and_target_probs_equal: bool
499
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
500
- draft_probs, target_probs = (F.softmax(
501
- torch.rand(self.vocab_size, dtype=torch.float32),
502
- dim=-1,
503
- ) for _ in range(2))
504
-
505
- num_reference_probs = 100
506
- reference_probs = F.softmax(
507
- torch.rand(num_reference_probs,
508
- self.vocab_size,
509
- dtype=torch.float32),
510
- dim=-1,
511
- )
512
-
513
- if draft_and_target_probs_equal:
514
- target_probs = draft_probs.clone()
515
-
516
- return draft_probs, target_probs, reference_probs
517
-
518
- def run_and_compare_distributions(self, draft_probs: torch.Tensor,
519
- target_probs: torch.Tensor,
520
- reference_probs: torch.Tensor,
521
- num_samples: int) -> tuple[float, float]:
522
- # Sample using rejection sampling.
523
- rej_sample_probs = self._estimate_rejection_sampling_pdf(
524
- draft_probs, target_probs, num_samples)
525
-
526
- # Average distance from reference probs.
527
- reference_vs_rejsample_dist = torch.dist(
528
- reference_probs,
529
- rej_sample_probs).item() / reference_probs.shape[0]
530
- target_vs_rejsample_dist = torch.dist(target_probs,
531
- rej_sample_probs).item()
532
-
533
- return reference_vs_rejsample_dist, target_vs_rejsample_dist
534
-
535
- def _estimate_rejection_sampling_pdf(
536
- self,
537
- draft_probs: torch.Tensor,
538
- target_probs: torch.Tensor,
539
- num_samples: int,
540
- ) -> torch.Tensor:
541
- # Repeat draft probs num_samples times.
542
- draft_probs = draft_probs.reshape(1, self.k, self.vocab_size).repeat(
543
- num_samples, 1, 1)
544
-
545
- # Repeat target probs num_samples * (k + 1) times.
546
- # Rejection sampler requires bonus token probs, but they aren't used.
547
- target_probs = target_probs.reshape(1, 1, self.vocab_size).repeat(
548
- num_samples, self.k + 1, 1)
549
-
550
- # Randomly sample draft token ids from draft probs.
551
- draft_token_ids = torch.multinomial(draft_probs[:, 0, :],
552
- num_samples=1,
553
- replacement=True).reshape(
554
- num_samples, self.k)
555
-
556
- # Bonus tokens not used but required.
557
- bonus_token_ids = torch.zeros((1, self.num_bonus_tokens),
558
- dtype=torch.int64,
559
- device="cuda").repeat(num_samples, 1)
560
-
561
- # Get output tokens via rejection sampling.
562
- output_token_ids = self.rejection_sampler(target_probs.to("cuda"),
563
- bonus_token_ids.to("cuda"),
564
- draft_probs.to("cuda"),
565
- draft_token_ids.to("cuda"))
566
-
567
- # Remove bonus tokens
568
- output_token_ids = output_token_ids[:, :-1].flatten()
569
-
570
- # Estimate probability density function
571
- hist = torch.histogram(output_token_ids.to(dtype=torch.float,
572
- device="cpu"),
573
- bins=self.vocab_size,
574
- range=self.vocab_range,
575
- density=True)
576
-
577
- return hist.hist
tests/samplers/test_typical_acceptance_sampler.py
deleted
-480
@@ -1,480 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests for rejection sampling."""
4
-
5
-import pytest
6
-import torch
7
-
8
-from vllm.model_executor.layers.typical_acceptance_sampler import (
9
- TypicalAcceptanceSampler)
10
-from vllm.model_executor.utils import set_random_seed
11
-
12
-CUDA_DEVICES = [f"cuda:{i}" for i in range(1)]
13
-
14
-
15
-@pytest.fixture(scope="function", autouse=True)
16
-def use_v0_only(monkeypatch):
17
- """
18
- This file tests V0 internals, so set VLLM_USE_V1=0.
19
- """
20
- monkeypatch.setenv('VLLM_USE_V1', '0')
21
-
22
-
23
-def get_zero_temperature_prob_dist(batch_size, k, vocab_size):
24
- """
25
- Generates a fake temperature zero probability distribution.
26
- Returns:
27
- 1. A fake temperature zero probability distribution of shape
28
- [batch_size, k, vocab_size]
29
- 2. Tensor of shape [batch_size, k] containing the token ids
30
- of the probability 1.0 tokens at each position.
31
- """
32
- # Simulate temperature 0 probability distribution for target probabilities
33
- # and create target probabilities such that only 1 token id has
34
- # probability 1.0
35
- target_probs = torch.rand(batch_size, k, vocab_size, dtype=torch.float32)
36
- probs = torch.rand(batch_size, k, vocab_size)
37
- _, zero_temperature_token_ids = torch.max(probs, dim=-1)
38
- # set the probability of the tokens with ids in zero_temperature_token_ids
39
- # to 1 and the rest to 0.
40
- target_probs = torch.zeros_like(probs).scatter_(
41
- -1, zero_temperature_token_ids.unsqueeze(-1), 1.0)
42
- return target_probs, zero_temperature_token_ids
43
-
44
-
45
-def get_draft_token_ids(batch_size: int, k: int, vocab_size: int,
46
- token_ids_to_exclude: torch.Tensor):
47
- """
48
- Returns a tensor of shape [batch_size, k] of fake draft token ids
49
- drawn randomly from a vocab of size vocab_size. We however ensure
50
- that token_ids from token_ids_to_exclude are excluded at the
51
- corresponding positions.
52
- """
53
- draft_token_ids = torch.empty(batch_size, k, dtype=torch.long)
54
- for i in range(batch_size):
55
- for j in range(k):
56
- # Generate a random token ID excluding token_ids_to_exclude[i, j]
57
- while True:
58
- token_id = torch.randint(0, vocab_size, (1, )).item()
59
- if token_id != token_ids_to_exclude[i, j]:
60
- draft_token_ids[i, j] = token_id
61
- break
62
- return draft_token_ids
63
-
64
-
65
-def get_acceptance_sampler(
66
- posterior_threshold: float = 0.03,
67
- posterior_alpha: float = 0.9,
68
- strict_mode: bool = False,
69
-) -> TypicalAcceptanceSampler:
70
- """
71
- Initializes and returns a TypicalAcceptanceSampler.
72
- """
73
- return TypicalAcceptanceSampler(posterior_threshold, posterior_alpha,
74
- strict_mode)
75
-
76
-
77
-@pytest.mark.parametrize("k", list(range(1, 6)))
78
-@pytest.mark.parametrize("vocab_size", [30_000, 50_000])
79
-@pytest.mark.parametrize("batch_size", list(range(1, 32)))
80
-@pytest.mark.parametrize("device", CUDA_DEVICES)
81
-@torch.inference_mode()
82
-def test_no_crash_with_varying_dims(k: int, vocab_size: int, batch_size: int,
83
- device: str):
84
- """
85
- Tests that the TypicalAcceptancSampler forward succeeds for
86
- different combinations of k, vocab_size, batch_size and num devices.
87
- """
88
- torch.set_default_device(device)
89
- typical_acceptance_sampler = get_acceptance_sampler()
90
- typical_acceptance_sampler.init_gpu_tensors(device=device)
91
- target_with_bonus_probs = torch.rand(batch_size,
92
- k + 1,
93
- vocab_size,
94
- dtype=torch.float32)
95
- bonus_token_ids = torch.randint(low=0,
96
- high=vocab_size,
97
- size=(batch_size, 1),
98
- dtype=torch.int64)
99
- draft_token_ids = torch.randint(low=0,
100
- high=vocab_size,
101
- size=(batch_size, k),
102
- dtype=torch.int64)
103
- # Verify that sampling succeeds for all cases.
104
- typical_acceptance_sampler(target_with_bonus_probs,
105
- bonus_token_ids,
106
- draft_probs=None,
107
- draft_token_ids=draft_token_ids)
108
-
109
-
110
-@pytest.mark.parametrize("above_or_below_vocab_range", ["above", "below"])
111
-@pytest.mark.parametrize("which_token_ids",
112
- ["bonus_token_ids", "draft_token_ids"])
113
-@pytest.mark.parametrize("device", CUDA_DEVICES)
114
-@torch.inference_mode()
115
-def test_raises_when_vocab_oob(above_or_below_vocab_range: str,
116
- which_token_ids: str, device: str):
117
- """
118
- Tests that we throw an exception of the token ids fall outside
119
- the bound of the provided vocabulary.
120
- """
121
- k = 3
122
- batch_size = 5
123
- vocab_size = 30_000
124
- torch.set_default_device(device)
125
- typical_acceptance_sampler = get_acceptance_sampler(strict_mode=True)
126
- typical_acceptance_sampler.init_gpu_tensors(device=device)
127
- target_with_bonus_probs = torch.rand(batch_size,
128
- k + 1,
129
- vocab_size,
130
- dtype=torch.float32)
131
- bonus_token_ids = torch.randint(low=0,
132
- high=vocab_size,
133
- size=(batch_size, 1),
134
- dtype=torch.int64)
135
- draft_token_ids = torch.randint(low=0,
136
- high=vocab_size,
137
- size=(batch_size, k),
138
- dtype=torch.int64)
139
- # Verify that appropriate exceptions are thrown for out
140
- # of bound vocabs.
141
- oob_token_ids = None
142
- if which_token_ids == "bonus_token_ids":
143
- oob_token_ids = bonus_token_ids
144
- elif which_token_ids == "draft_token_ids":
145
- oob_token_ids = draft_token_ids
146
- else:
147
- raise AssertionError()
148
-
149
- if above_or_below_vocab_range == "above":
150
- rogue_token_id = vocab_size + 1
151
- elif above_or_below_vocab_range == "below":
152
- rogue_token_id = -1
153
- else:
154
- raise AssertionError()
155
-
156
- oob_token_ids[0][0] = rogue_token_id
157
-
158
- with pytest.raises(AssertionError):
159
- typical_acceptance_sampler(target_with_bonus_probs,
160
- bonus_token_ids,
161
- draft_probs=None,
162
- draft_token_ids=draft_token_ids)
163
-
164
-
165
-@pytest.mark.parametrize("seed", list(range(10)))
166
-@pytest.mark.parametrize("device", CUDA_DEVICES)
167
-@torch.inference_mode()
168
-def test_uniform_target_distribution_accepts_all_tokens(
169
- seed: int, device: str):
170
- """
171
- Test the TypicalAcceptanceSampler with a uniform target probability
172
- distribution.
173
-
174
- This test verifies that when provided with a uniform target probability
175
- distribution, the TypicalAcceptanceSampler accepts all draft tokens. The
176
- entropy of the uniform target distribution being high should lead to all
177
- draft tokens being accepted.
178
- """
179
- set_random_seed(seed)
180
- k = 3
181
- batch_size = 5
182
- vocab_size = 30_000
183
- torch.set_default_device(device)
184
- typical_acceptance_sampler = get_acceptance_sampler(strict_mode=True)
185
- typical_acceptance_sampler.init_gpu_tensors(device=device)
186
- target_with_bonus_probs = torch.rand(batch_size,
187
- k + 1,
188
- vocab_size,
189
- dtype=torch.float32)
190
- draft_token_ids = torch.randint(low=0,
191
- high=vocab_size,
192
- size=(batch_size, k),
193
- dtype=torch.int64)
194
- bonus_token_ids = torch.randint(low=0,
195
- high=vocab_size,
196
- size=(batch_size, 1),
197
- dtype=torch.int64)
198
- output_token_ids = typical_acceptance_sampler(
199
- target_with_bonus_probs,
200
- bonus_token_ids,
201
- draft_probs=None,
202
- draft_token_ids=draft_token_ids)
203
- # We are using a uniform target probability distribution.
204
- # For a uniform distribution the entropy is very high and it
205
- # should lead to all draft tokens being accepted. Verify that.
206
- assert output_token_ids.shape[0] == batch_size
207
- assert output_token_ids.shape[1] == (k + 1)
208
- assert torch.all(output_token_ids[:, -1] == bonus_token_ids.squeeze())
209
-
210
- assert torch.all(output_token_ids[:, :k] == draft_token_ids)
211
-
212
-
213
-@pytest.mark.parametrize("seed", list(range(10)))
214
-@pytest.mark.parametrize("device", CUDA_DEVICES)
215
-@torch.inference_mode()
216
-def test_temperature_zero_target_distribution(seed: int, device: str):
217
- """
218
- Test the TypicalAcceptanceSampler with a zero-temperature target
219
- probability distribution.
220
-
221
- This test verifies that when using a zero-temperature target probability
222
- distribution, where only one token has a probability of 1.0, the
223
- TypicalAcceptanceSampler correctly rejects all draft tokens that do not
224
- match this probability. Additionally, it ensures that when all draft
225
- tokens are rejected, the sampler falls back to greedy sampling to select a
226
- single token from the target distribution.
227
- """
228
- set_random_seed(seed)
229
- k = 3
230
- batch_size = 5
231
- vocab_size = 30_000
232
- torch.set_default_device(device)
233
-
234
- typical_acceptance_sampler = get_acceptance_sampler(strict_mode=True)
235
- typical_acceptance_sampler.init_gpu_tensors(device=device)
236
- # Simulate temperature 0 probability distribution for target probabilities
237
- # and create target probabilities such that only 1 token id has
238
- # probability 1.0
239
- target_with_bonus_probs, zero_temperature_token_ids = \
240
- get_zero_temperature_prob_dist(batch_size, k + 1, vocab_size)
241
- zero_temperature_token_ids = zero_temperature_token_ids[:, :-1]
242
- # Populate draft_token_ids such that they exclude the token_ids
243
- # with probability = 1.0
244
- draft_token_ids = get_draft_token_ids(batch_size, k, vocab_size,
245
- zero_temperature_token_ids)
246
- bonus_token_ids = torch.randint(low=0,
247
- high=vocab_size,
248
- size=(batch_size, 1),
249
- dtype=torch.int64)
250
- # The target probaility distribution is a temperature zero distribution
251
- # with zero entropy. Since our draft token ids don't match the probability
252
- # 1.0 tokens in the target distribution we will reject all of them and
253
- # fallback to the greedy sampling for selecting 1 token for each sequence.
254
- # Verify the same.
255
- output_token_ids = typical_acceptance_sampler(
256
- target_with_bonus_probs,
257
- bonus_token_ids,
258
- draft_probs=None,
259
- draft_token_ids=draft_token_ids)
260
- assert output_token_ids.shape[0] == batch_size
261
- assert output_token_ids.shape[1] == (k + 1)
262
- assert torch.all(output_token_ids[:, -1] == -1)
263
- assert torch.all(output_token_ids[:, 0] == zero_temperature_token_ids[:,
264
- 0])
265
-
266
-
267
-@pytest.mark.parametrize("seed", list(range(10)))
268
-@pytest.mark.parametrize("device", CUDA_DEVICES)
269
-@torch.inference_mode()
270
-def test_mixed_target_distribution(seed: int, device: str):
271
- """
272
- Test the TypicalAcceptanceSampler with a mixed target probability
273
- distribution.
274
-
275
- This test ensures that the TypicalAcceptanceSampler handles a mixed
276
- target probability distribution correctly. Specifically, it uses a
277
- zero-temperature distribution for some sequences and a uniform
278
- distribution for others. The test verifies that:
279
-
280
- - For sequences with a zero-temperature distribution, only the token
281
- with a probability of 1.0 is accepted, and all other tokens are rejected.
282
- - For sequences with a uniform distribution, all draft tokens are
283
- accepted.
284
- """
285
- set_random_seed(seed)
286
- k = 3
287
- batch_size = 4
288
- vocab_size = 30_000
289
- torch.set_default_device(device)
290
- typical_acceptance_sampler = get_acceptance_sampler(strict_mode=True)
291
- typical_acceptance_sampler.init_gpu_tensors(device=device)
292
- # For sequences 0 and 2 set the distribution to a temperature
293
- # zero distribution. For sequences 1 and 3 set it to a uniform
294
- # distribution.
295
- target_with_bonus_probs, zero_temperature_token_ids = \
296
- get_zero_temperature_prob_dist(batch_size, k + 1, vocab_size)
297
- zero_temperature_token_ids = zero_temperature_token_ids[:, :-1]
298
- target_probs = target_with_bonus_probs[:, :-1]
299
- draft_token_ids = get_draft_token_ids(batch_size, k, vocab_size,
300
- zero_temperature_token_ids)
301
- uniform_probs = torch.rand(2, k, vocab_size, dtype=torch.float32)
302
- target_probs[[1, 3]] = uniform_probs
303
- bonus_token_ids = torch.randint(low=0,
304
- high=vocab_size,
305
- size=(batch_size, 1),
306
- dtype=torch.int64)
307
- output_token_ids = typical_acceptance_sampler(
308
- target_with_bonus_probs,
309
- bonus_token_ids,
310
- draft_probs=None,
311
- draft_token_ids=draft_token_ids)
312
- # verify the shape of output_token_ids
313
- assert output_token_ids.shape[0] == batch_size
314
- assert output_token_ids.shape[1] == (k + 1)
315
- # For sequences 0 and 2 verify that only 1 token is accepted
316
- # which is the token with probability 1.0 in the target distribution
317
- # at position 0.
318
- assert torch.all(output_token_ids[[0, 2], 1:] == -1)
319
- assert (torch.all(output_token_ids[[0, 2],
320
- 0] == zero_temperature_token_ids[[0, 2],
321
- 0]))
322
- # For sequences 1 and 3 verify that all tokens are accepted since the
323
- # target probability distribution is uniform. In addition verify that
324
- # we also accept the bonus tokens.
325
- assert torch.all(
326
- output_token_ids[[1, 3], :-1] == draft_token_ids[[1, 3], :])
327
- assert torch.all(output_token_ids[[1, 3], -1] != -1)
328
-
329
-
330
-@pytest.mark.parametrize("seed", list(range(10)))
331
-@pytest.mark.parametrize("device", CUDA_DEVICES)
332
-@torch.inference_mode()
333
-def test_accept_tokens_partially(seed: int, device: str):
334
- """
335
- Test the TypicalAcceptanceSampler's behavior when only a subset of draft
336
- tokens should be accepted.
337
-
338
- This test verifies that the TypicalAcceptanceSampler correctly accepts or
339
- rejects draft tokens based on a zero-temperature target probability
340
- distribution. Specifically, it ensures that:
341
-
342
- - When all draft tokens match tokens with a probability of 1.0 in the
343
- target distribution, all draft tokens are accepted.
344
- - When only some draft tokens match tokens with a probability of 1.0 in
345
- the target distribution, only those matching tokens are accepted, and the
346
- rest are rejected.
347
- """
348
- set_random_seed(seed)
349
- k = 5
350
- batch_size = 1
351
- vocab_size = 30_000
352
- torch.set_default_device(device)
353
- typical_acceptance_sampler = get_acceptance_sampler(strict_mode=True)
354
- typical_acceptance_sampler.init_gpu_tensors(device=device)
355
- # Create a temperature zero target probability distribution and ensure
356
- # all draft token ids correspond to the tokens with 1.0 probability.
357
- # Verify that all of them are accepted.
358
- target_with_bonus_probs, zero_temperature_token_ids = \
359
- get_zero_temperature_prob_dist(batch_size, k + 1, vocab_size)
360
- zero_temperature_token_ids = zero_temperature_token_ids[:, :-1]
361
- draft_token_ids = zero_temperature_token_ids
362
- bonus_token_ids = torch.randint(low=0,
363
- high=vocab_size,
364
- size=(batch_size, 1),
365
- dtype=torch.int64)
366
- output_token_ids = typical_acceptance_sampler(
367
- target_with_bonus_probs,
368
- bonus_token_ids,
369
- draft_probs=None,
370
- draft_token_ids=draft_token_ids)
371
- assert output_token_ids.shape[0] == batch_size
372
- assert output_token_ids.shape[1] == (k + 1)
373
- assert torch.all(output_token_ids[:, 0:-1] == draft_token_ids)
374
- assert torch.all(output_token_ids[:, -1] == bonus_token_ids)
375
- # Next only keep the first 2 draft tokens same as the zero temperature
376
- # tokens. For the remaining 3 choose some other tokens. In the
377
- # response we will expect the first 2 tokens to be the same as the
378
- # draft tokens and the recovered token and rest as -1
379
- draft_token_ids_to_replace = get_draft_token_ids(
380
- batch_size, k, vocab_size, zero_temperature_token_ids)
381
- draft_token_ids = torch.cat(
382
- (draft_token_ids[:, :2], draft_token_ids_to_replace[:, -3:]), dim=1)
383
- output_token_ids = typical_acceptance_sampler(
384
- target_with_bonus_probs,
385
- bonus_token_ids,
386
- draft_probs=None,
387
- draft_token_ids=draft_token_ids)
388
- assert output_token_ids.shape[0] == batch_size
389
- assert output_token_ids.shape[1] == (k + 1)
390
- assert torch.all(output_token_ids[:, :2] == draft_token_ids[:, :2])
391
- assert torch.all(
392
- output_token_ids[:, 2] == target_with_bonus_probs.argmax(-1)[:, 2])
393
- assert torch.all(output_token_ids[:, -3:] == -1)
394
-
395
-
396
-@pytest.mark.parametrize("seed", list(range(1)))
397
-@pytest.mark.parametrize("device", CUDA_DEVICES)
398
-@torch.inference_mode()
399
-def test_accept_tokens_set_non_default_posteriors(seed: int, device: str):
400
- """
401
- Test the TypicalAcceptanceSampler with custom posterior thresholds and
402
- alpha values. This test verifies that by modifying the posterior
403
- thresholds and alpha values we can change the acceptance behavior of the
404
- sampler.
405
- """
406
- set_random_seed(seed)
407
- k = 5
408
- batch_size = 1
409
- vocab_size = 30_000
410
- torch.set_default_device(device)
411
- typical_acceptance_sampler = get_acceptance_sampler(strict_mode=True)
412
- typical_acceptance_sampler.init_gpu_tensors(device=device)
413
- # Simulate temperature 0 probability distribution for target
414
- # probabilities and create target probabilities such that only 1 token
415
- # id has probability 1.0 and others have a very low probability of
416
- # 0.00001. Populate draft_token_ids such that they exclude the token_ids
417
- # with probability = 1.0. Without any changes to the posterior thresholds
418
- # none of the draft tokens are accepted.
419
- target_probs, zero_temperature_token_ids = get_zero_temperature_prob_dist(
420
- batch_size, k + 1, vocab_size)
421
- zero_temperature_token_ids = zero_temperature_token_ids[:, :-1]
422
- target_probs[target_probs == 0] = 0.00001
423
- draft_token_ids = get_draft_token_ids(batch_size, k, vocab_size,
424
- zero_temperature_token_ids)
425
- bonus_token_ids = torch.randint(low=0,
426
- high=vocab_size,
427
- size=(batch_size, 1),
428
- dtype=torch.int64)
429
- output_token_ids = typical_acceptance_sampler(
430
- target_probs,
431
- bonus_token_ids,
432
- draft_probs=None,
433
- draft_token_ids=draft_token_ids)
434
- assert output_token_ids.shape[0] == batch_size
435
- assert output_token_ids.shape[1] == (k + 1)
436
- assert torch.all(output_token_ids[:, 1:-1] == -1)
437
-
438
- # Change the posterior threshold values to 0.0 so that we will
439
- # now accept even draft tokens with very low probability in the
440
- # target distribution. Simulate and verify the same.
441
- typical_acceptance_sampler = TypicalAcceptanceSampler(
442
- strict_mode=True, posterior_threshold=0.0, posterior_alpha=0.0)
443
- typical_acceptance_sampler.init_gpu_tensors(device=device)
444
- output_token_ids = typical_acceptance_sampler(
445
- target_probs,
446
- bonus_token_ids,
447
- draft_probs=None,
448
- draft_token_ids=draft_token_ids)
449
- assert output_token_ids.shape[0] == batch_size
450
- assert output_token_ids.shape[1] == (k + 1)
451
- assert torch.all(output_token_ids[:, 0:-1] == draft_token_ids)
452
- assert torch.all(output_token_ids[:, -1] == bonus_token_ids)
453
-
454
-
455
-@pytest.mark.parametrize("seed", list(range(10)))
456
-@pytest.mark.parametrize("device", CUDA_DEVICES)
457
-@torch.inference_mode()
458
-def test_get_recovered_token_ids(seed: int, device: str):
459
- """
460
- Test the TypicalAcceptanceSampler's method for generating
461
- replacement token IDs.
462
-
463
- This test verifies that the `_get_recovered_token_ids` method of the
464
- TypicalAcceptanceSampler correctly identifies the token IDs to be used
465
- as recovered token IDs based on the target probability distribution.
466
- Specifically, it ensures that the method correctly identifies the
467
- tokens with the highest probability for each sequence in the batch.
468
- """
469
- set_random_seed(seed)
470
- k = 10
471
- batch_size = 5
472
- vocab_size = 30_000
473
- torch.set_default_device(device)
474
- typical_acceptance_sampler = get_acceptance_sampler(strict_mode=True)
475
- typical_acceptance_sampler.init_gpu_tensors(device=device)
476
- target_probs = torch.rand(batch_size, k, vocab_size, dtype=torch.float32)
477
- expected_replacement_tokens = torch.argmax(target_probs, dim=-1)
478
- actual_replacement_tokens = (
479
- typical_acceptance_sampler._get_recovered_token_ids(target_probs))
480
- assert torch.all(expected_replacement_tokens == actual_replacement_tokens)
tests/spec_decode/__init__.py
tests/spec_decode/conftest.py
deleted
-12
@@ -1,12 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import pytest
4
-
5
-
6
-@pytest.fixture(scope="function", autouse=True)
7
-def use_v0_only(monkeypatch):
8
- """
9
- Since this module is V0 only, set VLLM_USE_V1=0 for
10
- all tests in the module.
11
- """
12
- monkeypatch.setenv('VLLM_USE_V1', '0')
tests/spec_decode/e2e/__init__.py
tests/spec_decode/e2e/conftest.py
deleted
-307
@@ -1,307 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from collections.abc import Sequence
5
-from itertools import cycle
6
-from typing import Optional, Union
7
-
8
-import pytest
9
-import torch
10
-
11
-from vllm import LLM, SamplingParams
12
-from vllm.distributed import cleanup_dist_env_and_memory
13
-from vllm.model_executor.utils import set_random_seed
14
-from vllm.sequence import PromptLogprobs, SampleLogprobs
15
-
16
-from ...models.utils import (TokensTextLogprobs,
17
- TokensTextLogprobsPromptLogprobs,
18
- check_logprobs_close, check_outputs_equal)
19
-from ...utils import RemoteOpenAIServer
20
-
21
-PROMPTS = [
22
- "Hello, my name is",
23
- "The president of the United States is",
24
- "The capital of France is",
25
- "The future of AI is",
26
- "San Francisco is know for its",
27
- "Facebook was created in 2004 by",
28
- "Curious George is a",
29
- "Python 3.11 brings improvements to its",
30
-]
31
-
32
-
33
-@pytest.fixture
34
-def test_llm_generator(common_llm_kwargs, per_test_common_llm_kwargs,
35
- test_llm_kwargs, seed):
36
-
37
- def generate():
38
- kwargs = {
39
- **common_llm_kwargs,
40
- **per_test_common_llm_kwargs,
41
- **test_llm_kwargs,
42
- }
43
-
44
- llm = LLM(**kwargs)
45
-
46
- if seed is not None:
47
- set_random_seed(seed)
48
-
49
- yield llm
50
-
51
- del llm
52
- cleanup_dist_env_and_memory()
53
-
54
- return generate
55
-
56
-
57
-def maybe_assert_ngram_worker(llm):
58
- # Verify the proposer worker is ngram if ngram is specified.
59
- if (llm.llm_engine.speculative_config is not None
60
- and llm.llm_engine.speculative_config.method == "ngram"):
61
- from vllm.spec_decode.ngram_worker import NGramWorker
62
- assert isinstance(
63
- llm.llm_engine.model_executor.driver_worker.proposer_worker,
64
- NGramWorker)
65
-
66
-
67
-def get_output_from_llm_generator(
68
- llm_generator, prompts,
69
- sampling_params) -> tuple[list[str], list[list[int]], float]:
70
- tokens: list[str] = []
71
- token_ids: list[list[int]] = []
72
- acceptance_rate: float = -1.0
73
- for llm in llm_generator():
74
- maybe_assert_ngram_worker(llm)
75
-
76
- outputs = llm.generate(prompts, sampling_params, use_tqdm=True)
77
-
78
- token_ids = [output.outputs[0].token_ids for output in outputs]
79
- tokens = [output.outputs[0].text for output in outputs]
80
-
81
- # Fetch acceptance rate if logging is enabled.
82
- if stat_loggers := getattr(llm.llm_engine, "stat_loggers", None):
83
- stat_logger = stat_loggers["prometheus"]
84
- acceptance_rate = (stat_logger.metrics.
85
- gauge_spec_decode_draft_acceptance_rate.labels(
86
- **stat_logger.labels)._value.get())
87
- del llm
88
-
89
- return tokens, token_ids, acceptance_rate
90
-
91
-
92
-def check_logprobs_correctness(
93
- spec_outputs: Sequence[Union[TokensTextLogprobs,
94
- TokensTextLogprobsPromptLogprobs]],
95
- baseline_outputs: Sequence[Union[TokensTextLogprobs,
96
- TokensTextLogprobsPromptLogprobs]],
97
- disable_logprobs: bool = False,
98
-):
99
- """Compare sampled and prompt logprobs between baseline and spec decoding
100
- """
101
- if not disable_logprobs:
102
- return check_logprobs_close(
103
- outputs_0_lst=baseline_outputs,
104
- outputs_1_lst=spec_outputs,
105
- name_0="org",
106
- name_1="sd",
107
- )
108
-
109
- # Check correctness when disable_logprobs == True
110
- for spec_output, baseline_output in zip(spec_outputs, baseline_outputs):
111
- # Check generated token logprobs.
112
- spec_logprobs = spec_output[2]
113
- baseline_logprobs = baseline_output[2]
114
- _check_logprobs_when_output_disabled(spec_logprobs,
115
- baseline_logprobs,
116
- is_prompt_logprobs=False)
117
-
118
- # Check prompt logprobs too, if they exist
119
- if len(baseline_output) == 4:
120
- assert len(spec_output) == 4
121
- spec_prompt_logprobs = spec_output[3]
122
- baseline_prompt_logprobs = baseline_output[3]
123
- _check_logprobs_when_output_disabled(spec_prompt_logprobs,
124
- baseline_prompt_logprobs,
125
- is_prompt_logprobs=True)
126
-
127
-
128
-def _check_logprobs_when_output_disabled(
129
- spec_logprobs: Union[Optional[PromptLogprobs], SampleLogprobs],
130
- baseline_logprobs: Union[Optional[PromptLogprobs], SampleLogprobs],
131
- is_prompt_logprobs: bool = False,
132
-):
133
- # Prompt logprobs are optional
134
- if is_prompt_logprobs and baseline_logprobs is None:
135
- assert spec_logprobs is None
136
- return
137
-
138
- assert spec_logprobs is not None
139
- assert baseline_logprobs is not None
140
- assert len(spec_logprobs) == len(baseline_logprobs)
141
-
142
- # For each generated position of the sequence.
143
- for pos, (spec_pos_logprobs, baseline_pos_logprobs) in enumerate(
144
- zip(spec_logprobs, baseline_logprobs)):
145
-
146
- # First prompt logprob is expected to be None
147
- if is_prompt_logprobs and baseline_pos_logprobs is None:
148
- assert spec_pos_logprobs is None
149
- assert pos == 0
150
- continue
151
-
152
- assert spec_pos_logprobs is not None
153
- assert baseline_pos_logprobs is not None
154
-
155
- # When disabled, the 1 logprob is returned with dummy values for the
156
- # score and rank, but the token id should match the baseline model
157
- assert len(spec_pos_logprobs) == 1
158
- (spec_pos_logprob_token_id,
159
- spec_pos_logprob) = next(iter(spec_pos_logprobs.items()))
160
- assert spec_pos_logprob.rank == -1
161
- assert spec_pos_logprob.logprob == 0.0
162
- if isinstance(spec_pos_logprob_token_id, torch.Tensor):
163
- spec_pos_logprob_token_id = spec_pos_logprob_token_id.item()
164
- assert spec_pos_logprob_token_id in baseline_pos_logprobs
165
-
166
-
167
-def run_equality_correctness_test(
168
- vllm_runner,
169
- common_llm_kwargs,
170
- per_test_common_llm_kwargs,
171
- baseline_llm_kwargs,
172
- test_llm_kwargs,
173
- batch_size: int,
174
- max_output_len: int,
175
- seed: Optional[int] = 0,
176
- temperature: float = 0.0,
177
- disable_seed: bool = False,
178
- ignore_eos: bool = True,
179
- ensure_all_accepted: bool = False,
180
- expected_acceptance_rate: Optional[float] = None,
181
- logprobs: Optional[int] = None,
182
- prompt_logprobs: Optional[int] = None,
183
- disable_logprobs: bool = False):
184
-
185
- org_args = {
186
- **common_llm_kwargs,
187
- **per_test_common_llm_kwargs,
188
- **baseline_llm_kwargs,
189
- }
190
-
191
- sd_args = {
192
- **common_llm_kwargs,
193
- **per_test_common_llm_kwargs,
194
- **test_llm_kwargs,
195
- }
196
-
197
- prompts = [prompt for prompt, _ in zip(cycle(PROMPTS), range(batch_size))]
198
-
199
- if disable_seed:
200
- seed = None
201
-
202
- sampling_params = SamplingParams(temperature=temperature,
203
- max_tokens=max_output_len,
204
- seed=seed,
205
- ignore_eos=ignore_eos,
206
- logprobs=logprobs,
207
- prompt_logprobs=prompt_logprobs)
208
-
209
- with vllm_runner(**org_args) as vllm_model:
210
- org_outputs = vllm_model.generate_w_logprobs(prompts, sampling_params)
211
-
212
- with vllm_runner(**sd_args) as vllm_model:
213
- if ensure_all_accepted or expected_acceptance_rate is not None:
214
- # Force log interval to be 0 to catch all metrics.
215
- stat_logger = vllm_model.model.llm_engine.stat_loggers[
216
- 'prometheus']
217
- stat_logger.local_interval = -100
218
-
219
- sd_outputs = vllm_model.generate_w_logprobs(prompts, sampling_params)
220
-
221
- if ensure_all_accepted or expected_acceptance_rate is not None:
222
- acceptance_rate = (stat_logger.metrics.
223
- gauge_spec_decode_draft_acceptance_rate.labels(
224
- **stat_logger.labels)._value.get())
225
-
226
- if ensure_all_accepted:
227
- assert True
228
- # FIXME: ci fails to log acceptance rate.
229
- # It works locally.
230
- # assert acceptance_rate == 1.0
231
-
232
- if expected_acceptance_rate is not None:
233
- assert acceptance_rate >= expected_acceptance_rate - 1e-2
234
-
235
- # Only pass token entries, not the logprobs
236
- check_outputs_equal(outputs_0_lst=[out[0:2] for out in org_outputs],
237
- outputs_1_lst=[out[0:2] for out in sd_outputs],
238
- name_0="org",
239
- name_1="sd")
240
-
241
- # Check logprobs if requested
242
- if logprobs is not None or prompt_logprobs is not None:
243
- check_logprobs_correctness(sd_outputs, org_outputs, disable_logprobs)
244
-
245
-
246
-def run_equality_correctness_test_tp(model,
247
- common_llm_kwargs,
248
- per_test_common_llm_kwargs,
249
- baseline_llm_kwargs,
250
- test_llm_kwargs,
251
- batch_size: int,
252
- max_output_len: int,
253
- seed: int = 0,
254
- temperature: float = 0.0,
255
- logprobs: Optional[int] = None):
256
- """Helper method that compares the outputs of both the baseline LLM and
257
- the test LLM. It asserts greedy equality, e.g. that the outputs are exactly
258
- the same when temperature is zero.
259
- """
260
- arg1 = common_llm_kwargs + per_test_common_llm_kwargs + baseline_llm_kwargs
261
- arg2 = common_llm_kwargs + per_test_common_llm_kwargs + test_llm_kwargs
262
- env1 = env2 = None
263
-
264
- max_wait_seconds = 240
265
- results = []
266
-
267
- prompts = [prompt for prompt, _ in zip(cycle(PROMPTS), range(batch_size))]
268
- for args, env in ((arg1, env1), (arg2, env2)):
269
- with RemoteOpenAIServer(model,
270
- args,
271
- env_dict=env,
272
- max_wait_seconds=max_wait_seconds) as server:
273
- client = server.get_client()
274
-
275
- completion = client.completions.create(model=model,
276
- prompt=prompts,
277
- max_tokens=max_output_len,
278
- seed=seed,
279
- temperature=temperature,
280
- logprobs=logprobs)
281
-
282
- results.append({
283
- "test":
284
- "seeded_sampling",
285
- "text": [choice.text for choice in completion.choices],
286
- "logprobs": [choice.logprobs for choice in completion.choices],
287
- "finish_reason":
288
- [choice.finish_reason for choice in completion.choices],
289
- "usage":
290
- completion.usage,
291
- })
292
-
293
- n = len(results) // 2
294
- arg1_results = results[:n]
295
- arg2_results = results[n:]
296
- # Separate logprobs to avoid asserting exact equality.
297
- arg1_logprobs = [r.pop("logprobs") for r in arg1_results]
298
- arg2_logprobs = [r.pop("logprobs") for r in arg2_results]
299
-
300
- for arg1_result, arg2_result in zip(arg1_results, arg2_results):
301
- assert arg1_result == arg2_result, (
302
- f"Results for {model=} are not the same with {arg1=} and {arg2=}. "
303
- f"{arg1_result=} != {arg2_result=}")
304
- if logprobs:
305
- for logs1, logs2 in zip(arg1_logprobs, arg2_logprobs):
306
- for l1, l2 in zip(logs1, logs2):
307
- assert l1.tokens == l2.tokens
tests/spec_decode/e2e/test_compatibility.py
deleted
-66
@@ -1,66 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import pytest
5
-
6
-from vllm import SamplingParams
7
-
8
-from .conftest import get_output_from_llm_generator
9
-
10
-
11
-@pytest.mark.parametrize("common_llm_kwargs",
12
- [{
13
- "model": "meta-llama/Llama-3.2-1B-Instruct",
14
- }])
15
-@pytest.mark.parametrize(
16
- "per_test_common_llm_kwargs",
17
- [
18
- {
19
- # Speculative max model len > overridden max model len should raise.
20
- "speculative_config": {
21
- "model": "JackFram/llama-68m",
22
- "num_speculative_tokens": 5,
23
- "max_model_len": 129,
24
- },
25
- "max_model_len": 128,
26
- },
27
- {
28
- # Speculative max model len > draft max model len should raise.
29
- # https://huggingface.co/JackFram/llama-68m/blob/3b606af5198a0b26762d589a3ee3d26ee6fa6c85/config.json#L12
30
- "speculative_config": {
31
- "model": "JackFram/llama-68m",
32
- "num_speculative_tokens": 5,
33
- "max_model_len": 2048 + 1,
34
- },
35
- },
36
- {
37
- # Speculative max model len > target max model len should raise.
38
- # https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct/blob/9213176726f574b556790deb65791e0c5aa438b6/config.json#L18
39
- "speculative_config": {
40
- "model": "JackFram/llama-68m",
41
- "num_speculative_tokens": 5,
42
- "max_model_len": 131072 + 1,
43
- },
44
- },
45
- ])
46
-@pytest.mark.parametrize("test_llm_kwargs", [{}])
47
-@pytest.mark.parametrize("seed", [1])
48
-def test_spec_decode_xfail_spec_max_model_len(test_llm_generator):
49
- """Verify that speculative decoding validates speculative_max_model_len.
50
- """
51
- output_len = 128
52
- temperature = 0.0
53
-
54
- prompts = [
55
- "Hello, my name is",
56
- ]
57
-
58
- sampling_params = SamplingParams(
59
- max_tokens=output_len,
60
- ignore_eos=True,
61
- temperature=temperature,
62
- )
63
-
64
- with pytest.raises(ValueError, match="cannot be larger than"):
65
- get_output_from_llm_generator(test_llm_generator, prompts,
66
- sampling_params)
tests/spec_decode/e2e/test_eagle_correctness.py
deleted
-480
@@ -1,480 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""This docstring details important information on the testing methodology.
4
-
5
-Most of the tests rely on "greedy equality", where we expect the output of
6
-speculative decoding on a sequence to exactly match the output of normal non-
7
-speculative decoding.
8
-
9
-Since speculative decoding with rejection sampling guarantees that the output
10
-distribution matches the target model's output distribution (up to hardware
11
-numerics, see https://arxiv.org/pdf/2302.01318.pdf), we can expect greedy
12
-equality.
13
-
14
-However, we still need to verify below scenario could be passed:
15
- * Batch size 1 greedy equality
16
- * Batch size >1 greedy equality
17
- * Test greedy equality under preemption
18
- * Test greedy equality under various number of speculative tokens.
19
-
20
-With those tests, we can say at least, EAGLE would not break the
21
-correctness for the target model outputs.
22
-"""
23
-
24
-import pytest
25
-
26
-from .conftest import run_equality_correctness_test
27
-
28
-# main model
29
-MAIN_MODEL = "JackFram/llama-68m"
30
-
31
-# speculative model
32
-SPEC_MODEL = "abhigoyal/vllm-eagle-llama-68m-random"
33
-
34
-# max. number of speculative tokens: this corresponds to
35
-# num_heads in the config.json of the speculator model.
36
-MAX_SPEC_TOKENS = 4
37
-
38
-# precision
39
-PRECISION = "float32"
40
-
41
-
42
-@pytest.mark.parametrize(
43
- "common_llm_kwargs",
44
- [{
45
- # Skip cuda graph recording for fast test.
46
- "enforce_eager": True,
47
-
48
- # Print spec metrics.
49
- "disable_log_stats": False,
50
-
51
- # Precision
52
- "dtype": PRECISION,
53
-
54
- # Main model
55
- "model_name": MAIN_MODEL,
56
- }])
57
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
58
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
59
-@pytest.mark.parametrize("test_llm_kwargs", [
60
- {
61
- "speculative_config": {
62
- "model": SPEC_MODEL,
63
- "num_speculative_tokens": MAX_SPEC_TOKENS,
64
- },
65
- },
66
-])
67
-@pytest.mark.parametrize("output_len", [
68
- 128,
69
-])
70
-@pytest.mark.parametrize("batch_size", [1, 32])
71
-@pytest.mark.parametrize("seed", [1])
72
-def test_eagle_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
73
- per_test_common_llm_kwargs,
74
- baseline_llm_kwargs, test_llm_kwargs,
75
- batch_size: int, output_len: int,
76
- seed: int):
77
-
78
- run_equality_correctness_test(vllm_runner, common_llm_kwargs,
79
- per_test_common_llm_kwargs,
80
- baseline_llm_kwargs, test_llm_kwargs,
81
- batch_size, output_len, seed)
82
-
83
-
84
-@pytest.mark.parametrize(
85
- "common_llm_kwargs",
86
- [{
87
- # Skip cuda graph recording for fast test.
88
- "enforce_eager": True,
89
-
90
- # Print spec metrics.
91
- "disable_log_stats": False,
92
-
93
- # Precision
94
- "dtype": PRECISION,
95
-
96
- # Main model
97
- "model_name": MAIN_MODEL,
98
- }])
99
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
100
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
101
-@pytest.mark.parametrize("test_llm_kwargs", [{
102
- "speculative_config": {
103
- "model": SPEC_MODEL,
104
- "num_speculative_tokens": MAX_SPEC_TOKENS,
105
- "disable_logprobs": False,
106
- },
107
-}, {
108
- "speculative_config": {
109
- "model": SPEC_MODEL,
110
- "num_speculative_tokens": MAX_SPEC_TOKENS,
111
- "disable_logprobs": True,
112
- },
113
-}])
114
-@pytest.mark.parametrize("output_len", [
115
- 128,
116
-])
117
-@pytest.mark.parametrize("batch_size", [8])
118
-@pytest.mark.parametrize("seed", [1])
119
-@pytest.mark.parametrize("logprobs", [1, 6])
120
-def test_eagle_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
121
- per_test_common_llm_kwargs,
122
- baseline_llm_kwargs, test_llm_kwargs,
123
- batch_size: int, output_len: int, seed: int,
124
- logprobs: int):
125
-
126
- run_equality_correctness_test(
127
- vllm_runner,
128
- common_llm_kwargs,
129
- per_test_common_llm_kwargs,
130
- baseline_llm_kwargs,
131
- test_llm_kwargs,
132
- batch_size,
133
- output_len,
134
- seed,
135
- logprobs=logprobs,
136
- prompt_logprobs=logprobs,
137
- disable_logprobs=test_llm_kwargs["speculative_config"]
138
- ["disable_logprobs"])
139
-
140
-
141
-@pytest.mark.parametrize(
142
- "common_llm_kwargs",
143
- [{
144
- "enforce_eager": False,
145
-
146
- # Print spec metrics.
147
- "disable_log_stats": False,
148
-
149
- # Precision
150
- "dtype": PRECISION,
151
-
152
- # Main model
153
- "model_name": MAIN_MODEL,
154
- }])
155
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
156
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
157
-@pytest.mark.parametrize("test_llm_kwargs", [
158
- {
159
- "speculative_config": {
160
- "model": SPEC_MODEL,
161
- "num_speculative_tokens": MAX_SPEC_TOKENS,
162
- },
163
- },
164
-])
165
-@pytest.mark.parametrize("output_len", [
166
- 128,
167
-])
168
-@pytest.mark.parametrize("batch_size", [1, 32])
169
-@pytest.mark.parametrize("seed", [1])
170
-def test_eagle_e2e_greedy_correctness_cuda_graph(
171
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
172
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
173
- seed: int):
174
- """Verify greedy equality with cuda graph enabled and different
175
- batch sizes."""
176
- run_equality_correctness_test(vllm_runner, common_llm_kwargs,
177
- per_test_common_llm_kwargs,
178
- baseline_llm_kwargs, test_llm_kwargs,
179
- batch_size, output_len, seed)
180
-
181
-
182
-@pytest.mark.parametrize(
183
- "common_llm_kwargs",
184
- [{
185
- "block_size": 8,
186
- # 2 for small prompt, 256//8 for generated.
187
- "num_gpu_blocks_override": 2 + 256 // 8,
188
- "max_model_len": (2 + 256 // 8) * 8,
189
-
190
- # Skip cuda graph recording for fast test.
191
- "enforce_eager": True,
192
-
193
- # Precision
194
- "dtype": PRECISION,
195
-
196
- # Main model
197
- "model_name": MAIN_MODEL,
198
- }])
199
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
200
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
201
-@pytest.mark.parametrize("test_llm_kwargs", [
202
- {
203
- "speculative_config": {
204
- "model": SPEC_MODEL,
205
- "num_speculative_tokens": MAX_SPEC_TOKENS,
206
- },
207
- },
208
-])
209
-@pytest.mark.parametrize(
210
- "output_len",
211
- [
212
- # Use small output len for fast test.
213
- 128,
214
- ])
215
-@pytest.mark.parametrize("batch_size", [4])
216
-@pytest.mark.parametrize("seed", [1])
217
-def test_eagle_e2e_greedy_correctness_with_preemption(
218
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
219
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
220
- seed: int):
221
- """Verify greedy equality, even when some sequences are preempted mid-
222
- generation.
223
- """
224
- run_equality_correctness_test(vllm_runner, common_llm_kwargs,
225
- per_test_common_llm_kwargs,
226
- baseline_llm_kwargs, test_llm_kwargs,
227
- batch_size, output_len, seed)
228
-
229
-
230
-@pytest.mark.parametrize(
231
- "common_llm_kwargs",
232
- [{
233
- # Skip cuda graph recording for fast test.
234
- "enforce_eager": True,
235
-
236
- # Precision
237
- "dtype": PRECISION,
238
-
239
- # Main model
240
- "model_name": MAIN_MODEL,
241
- }])
242
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
243
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
244
-@pytest.mark.parametrize(
245
- "test_llm_kwargs",
246
- [
247
- {
248
- "speculative_config": {
249
- "model": SPEC_MODEL,
250
- "num_speculative_tokens": k,
251
- },
252
- }
253
- # Try a range of num. speculative tokens
254
- for k in range(1, 1 + MAX_SPEC_TOKENS)
255
- ])
256
-@pytest.mark.parametrize("batch_size", [2])
257
-@pytest.mark.parametrize(
258
- "output_len",
259
- [
260
- # Use smaller output len for fast test.
261
- 32,
262
- ])
263
-@pytest.mark.parametrize("seed", [1])
264
-def test_eagle_different_k(vllm_runner, common_llm_kwargs,
265
- per_test_common_llm_kwargs, baseline_llm_kwargs,
266
- test_llm_kwargs, batch_size: int, output_len: int,
267
- seed: int):
268
- """Verify that eagle speculative decoding produces exact equality
269
- to without spec decode with different values of num_speculative_tokens.
270
- """
271
- run_equality_correctness_test(vllm_runner, common_llm_kwargs,
272
- per_test_common_llm_kwargs,
273
- baseline_llm_kwargs, test_llm_kwargs,
274
- batch_size, output_len, seed)
275
-
276
-
277
-@pytest.mark.parametrize(
278
- "common_llm_kwargs",
279
- [{
280
- # Skip cuda graph recording for fast test.
281
- "enforce_eager": True,
282
-
283
- # Precision
284
- "dtype": PRECISION,
285
-
286
- # Main model
287
- "model_name": MAIN_MODEL,
288
- }])
289
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
290
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
291
-@pytest.mark.parametrize("test_llm_kwargs", [{
292
- "speculative_config": {
293
- "model": SPEC_MODEL,
294
- "num_speculative_tokens": MAX_SPEC_TOKENS,
295
- "disable_by_batch_size": 4,
296
- },
297
-}])
298
-@pytest.mark.parametrize("batch_size", [1, 5])
299
-@pytest.mark.parametrize(
300
- "output_len",
301
- [
302
- # Use smaller output len for fast test.
303
- 32,
304
- ])
305
-@pytest.mark.parametrize("seed", [1])
306
-def test_eagle_disable_queue(vllm_runner, common_llm_kwargs,
307
- per_test_common_llm_kwargs, baseline_llm_kwargs,
308
- test_llm_kwargs, batch_size: int, output_len: int,
309
- seed: int):
310
- """Verify that eagle speculative decoding produces exact equality
311
- to without spec decode when speculation is disabled for large
312
- batch sizes.
313
- """
314
- run_equality_correctness_test(vllm_runner, common_llm_kwargs,
315
- per_test_common_llm_kwargs,
316
- baseline_llm_kwargs, test_llm_kwargs,
317
- batch_size, output_len, seed)
318
-
319
-
320
-@pytest.mark.parametrize(
321
- "common_llm_kwargs",
322
- [{
323
- # Skip cuda graph recording for fast test.
324
- "enforce_eager": True,
325
-
326
- # Print spec metrics.
327
- "disable_log_stats": False,
328
-
329
- # Precision
330
- "dtype": "float16",
331
-
332
- # Main model
333
- "model_name": "meta-llama/Llama-2-7b-chat-hf",
334
- }])
335
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
336
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
337
-@pytest.mark.parametrize("test_llm_kwargs", [
338
- {
339
- "speculative_config": {
340
- "model": "yuhuili/EAGLE-llama2-chat-7B",
341
- "num_speculative_tokens": MAX_SPEC_TOKENS,
342
- },
343
- },
344
-])
345
-@pytest.mark.parametrize(
346
- "output_len",
347
- [
348
- # Use smaller output len for fast test.
349
- 32,
350
- ])
351
-@pytest.mark.parametrize("batch_size", [1, 5])
352
-@pytest.mark.parametrize("seed", [1])
353
-def test_llama2_eagle_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
354
- per_test_common_llm_kwargs,
355
- baseline_llm_kwargs,
356
- test_llm_kwargs, batch_size: int,
357
- output_len: int, seed: int):
358
-
359
- run_equality_correctness_test(vllm_runner,
360
- common_llm_kwargs,
361
- per_test_common_llm_kwargs,
362
- baseline_llm_kwargs,
363
- test_llm_kwargs,
364
- batch_size,
365
- output_len,
366
- seed,
367
- temperature=0.0)
368
-
369
-
370
-@pytest.mark.parametrize(
371
- "common_llm_kwargs",
372
- [{
373
- # 2 for small prompt, 256//16 for generated.
374
- "num_gpu_blocks_override": 2 + 256 // 16,
375
- "max_model_len": (2 + 256 // 16) * 16,
376
-
377
- # Skip cuda graph recording for fast test.
378
- "enforce_eager": True,
379
-
380
- # Print spec metrics.
381
- "disable_log_stats": False,
382
-
383
- # Precision
384
- "dtype": "float16",
385
-
386
- # Main model
387
- "model_name": "meta-llama/Meta-Llama-3-8B-Instruct",
388
- }])
389
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
390
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
391
-@pytest.mark.parametrize("test_llm_kwargs", [
392
- {
393
- "speculative_config": {
394
- "model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
395
- "num_speculative_tokens": MAX_SPEC_TOKENS,
396
- },
397
- },
398
-])
399
-@pytest.mark.parametrize(
400
- "output_len",
401
- [
402
- # Use smaller output len for fast test.
403
- 32,
404
- ])
405
-@pytest.mark.parametrize("batch_size", [1, 5])
406
-@pytest.mark.parametrize("seed", [1])
407
-def test_llama3_eagle_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
408
- per_test_common_llm_kwargs,
409
- baseline_llm_kwargs,
410
- test_llm_kwargs, batch_size: int,
411
- output_len: int, seed: int):
412
-
413
- run_equality_correctness_test(vllm_runner,
414
- common_llm_kwargs,
415
- per_test_common_llm_kwargs,
416
- baseline_llm_kwargs,
417
- test_llm_kwargs,
418
- batch_size,
419
- output_len,
420
- seed,
421
- temperature=0.0)
422
-
423
-
424
-@pytest.mark.parametrize(
425
- "common_llm_kwargs",
426
- [{
427
- # 2 for small prompt, 256//16 for generated.
428
- "num_gpu_blocks_override": 2 + 256 // 16,
429
- "max_model_len": (2 + 256 // 16) * 16,
430
-
431
- # Skip cuda graph recording for fast test.
432
- "enforce_eager": True,
433
-
434
- # Print spec metrics.
435
- "disable_log_stats": False,
436
-
437
- # Precision
438
- "dtype": "float16",
439
-
440
- # Main model
441
- "model_name": "Qwen/Qwen2-7B-Instruct",
442
- }])
443
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
444
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
445
-@pytest.mark.parametrize("test_llm_kwargs", [
446
- {
447
- "speculative_config": {
448
- "model": "yuhuili/EAGLE-Qwen2-7B-Instruct",
449
- "num_speculative_tokens": MAX_SPEC_TOKENS,
450
- },
451
- },
452
-])
453
-@pytest.mark.parametrize(
454
- "output_len",
455
- [
456
- # Use smaller output len for fast test.
457
- 32,
458
- ])
459
-@pytest.mark.parametrize("batch_size", [1, 5])
460
-@pytest.mark.parametrize("seed", [1])
461
-def test_qwen2_eagle_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
462
- per_test_common_llm_kwargs,
463
- baseline_llm_kwargs,
464
- test_llm_kwargs, batch_size: int,
465
- output_len: int, seed: int):
466
-
467
- run_equality_correctness_test(vllm_runner,
468
- common_llm_kwargs,
469
- per_test_common_llm_kwargs,
470
- baseline_llm_kwargs,
471
- test_llm_kwargs,
472
- batch_size,
473
- output_len,
474
- seed,
475
- temperature=0.0)
476
-
477
-
478
-if __name__ == "__main__":
479
- import pytest
480
- pytest.main([__file__])
tests/spec_decode/e2e/test_integration.py
deleted
-161
@@ -1,161 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests which cover integration of the speculative decoding framework with
4
-other features, e.g. cuda graphs.
5
-"""
6
-
7
-import pytest
8
-
9
-from .conftest import run_equality_correctness_test
10
-
11
-MAIN_MODEL = "JackFram/llama-68m"
12
-
13
-
14
-@pytest.mark.parametrize(
15
- "common_llm_kwargs",
16
- [{
17
- "model_name": "JackFram/llama-68m",
18
-
19
- # Verify equality when cuda graphs allowed.
20
- "enforce_eager": False,
21
-
22
- # The original model is float32, keep it for numerical stability.
23
- "dtype": "float32",
24
- }])
25
-@pytest.mark.parametrize(
26
- "per_test_common_llm_kwargs",
27
- [
28
- {
29
- # Identical models.
30
- "speculative_config": {
31
- "model": "JackFram/llama-68m",
32
- "num_speculative_tokens": 5,
33
- },
34
- },
35
- ])
36
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
37
-@pytest.mark.parametrize("test_llm_kwargs", [{}])
38
-@pytest.mark.parametrize("batch_size", [8])
39
-@pytest.mark.parametrize("output_len", [32])
40
-@pytest.mark.parametrize("seed", [1])
41
-def test_spec_decode_cuda_graph(vllm_runner, common_llm_kwargs,
42
- per_test_common_llm_kwargs,
43
- baseline_llm_kwargs, test_llm_kwargs,
44
- batch_size: int, output_len: int, seed: int):
45
- """Verify spec decode equality when cuda graphs are enabled.
46
- """
47
- run_equality_correctness_test(vllm_runner,
48
- common_llm_kwargs,
49
- per_test_common_llm_kwargs,
50
- baseline_llm_kwargs,
51
- test_llm_kwargs,
52
- batch_size,
53
- max_output_len=output_len,
54
- seed=seed,
55
- temperature=0.0)
56
-
57
-
58
-@pytest.mark.parametrize(
59
- "common_llm_kwargs",
60
- [{
61
- "model_name": "JackFram/llama-160m",
62
-
63
- # Skip cuda graph recording for fast test.
64
- "enforce_eager": True,
65
-
66
- # The original model is float32, keep it for numerical stability.
67
- "dtype": "float32",
68
- }])
69
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [])
70
-@pytest.mark.parametrize(
71
- "test_llm_kwargs",
72
- [
73
- # Explicitly specify draft model quantization
74
- {
75
- "speculative_config": {
76
- "model": "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit",
77
- "num_speculative_tokens": 5,
78
- "quantization": "gptq",
79
- },
80
- },
81
- # Explicitly specify GPTQ-based draft model to use marlin quantization
82
- {
83
- "speculative_config": {
84
- "model": "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit",
85
- "num_speculative_tokens": 5,
86
- "quantization": "marlin",
87
- },
88
- },
89
- # Not explicitly specify draft model quantization
90
- {
91
- "speculative_config": {
92
- "model": "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit",
93
- "num_speculative_tokens": 5,
94
- "quantization": None,
95
- },
96
- },
97
- ])
98
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
99
-@pytest.mark.parametrize("batch_size", [2])
100
-@pytest.mark.parametrize("seed", [1])
101
-def test_speculative_model_quantization_config(vllm_runner, common_llm_kwargs,
102
- per_test_common_llm_kwargs,
103
- baseline_llm_kwargs,
104
- test_llm_kwargs,
105
- batch_size: int, seed: int):
106
- """Verify spec decode works well with draft model quantization configs.
107
- """
108
- run_equality_correctness_test(vllm_runner,
109
- common_llm_kwargs,
110
- per_test_common_llm_kwargs,
111
- baseline_llm_kwargs,
112
- test_llm_kwargs,
113
- batch_size,
114
- max_output_len=32,
115
- seed=seed,
116
- temperature=0.0)
117
-
118
-
119
-@pytest.mark.parametrize(
120
- "common_llm_kwargs",
121
- [{
122
- "model_name": MAIN_MODEL,
123
-
124
- # Skip cuda graph recording for fast test.
125
- "enforce_eager": True,
126
-
127
- # The original model is float32, keep it for numerical stability.
128
- "dtype": "float32",
129
- }])
130
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
131
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
132
-@pytest.mark.parametrize("test_llm_kwargs", [{
133
- "speculative_config": {
134
- "model": "JackFram/llama-68m",
135
- "num_speculative_tokens": 3,
136
- "disable_mqa_scorer": True,
137
- },
138
-}])
139
-@pytest.mark.parametrize("batch_size", [1, 5])
140
-@pytest.mark.parametrize(
141
- "output_len",
142
- [
143
- # Use smaller output len for fast test.
144
- 32,
145
- ])
146
-@pytest.mark.parametrize("seed", [1])
147
-def test_mqa_scorer(vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
148
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int,
149
- output_len: int, seed: int):
150
- """Verify that speculative decoding generates the same output
151
- with batch expansion scorer and mqa scorer.
152
- """
153
- run_equality_correctness_test(vllm_runner,
154
- common_llm_kwargs,
155
- per_test_common_llm_kwargs,
156
- baseline_llm_kwargs,
157
- test_llm_kwargs,
158
- batch_size,
159
- max_output_len=output_len,
160
- seed=seed,
161
- temperature=0.0)
tests/spec_decode/e2e/test_integration_dist_tp2.py
deleted
-247
@@ -1,247 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests which cover integration of the speculative decoding framework with
4
-tensor parallelism.
5
-"""
6
-
7
-import json
8
-from typing import Optional
9
-
10
-import pytest
11
-import torch
12
-
13
-from vllm.platforms import current_platform
14
-
15
-from .conftest import run_equality_correctness_test_tp
16
-
17
-
18
-@pytest.mark.skipif(torch.cuda.device_count() < 2,
19
- reason="Need at least 2 GPUs to run the test.")
20
-@pytest.mark.parametrize(
21
- "common_llm_kwargs",
22
- [[
23
- # Skip cuda graph recording for fast test.
24
- "--enforce-eager",
25
- "--tensor-parallel-size",
26
- "2"
27
- ]])
28
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [[]])
29
-@pytest.mark.parametrize("baseline_llm_kwargs", [[]])
30
-@pytest.mark.parametrize("test_llm_kwargs", [
31
- [
32
- "--speculative_config",
33
- json.dumps({
34
- "model": "JackFram/llama-68m",
35
- "num_speculative_tokens": 3,
36
- }),
37
- ],
38
- [
39
- "--speculative_config",
40
- json.dumps({
41
- "model": "ngram",
42
- "num_speculative_tokens": 5,
43
- "prompt_lookup_max": 3,
44
- }),
45
- ],
46
-])
47
-@pytest.mark.parametrize("batch_size", [2])
48
-@pytest.mark.parametrize(
49
- "output_len",
50
- [
51
- # Use smaller output len for fast test.
52
- 32,
53
- ])
54
-@pytest.mark.parametrize("seed", [1])
55
-def test_target_model_tp_gt_1(common_llm_kwargs, per_test_common_llm_kwargs,
56
- baseline_llm_kwargs, test_llm_kwargs,
57
- batch_size: int, output_len: int, seed: int):
58
- """Verify greedy equality when tensor parallelism is used.
59
- """
60
- if current_platform.is_rocm():
61
- pytest.skip("hip is not well-supported yet")
62
- run_equality_correctness_test_tp("JackFram/llama-68m",
63
- common_llm_kwargs,
64
- per_test_common_llm_kwargs,
65
- baseline_llm_kwargs,
66
- test_llm_kwargs,
67
- batch_size,
68
- output_len,
69
- seed,
70
- temperature=0.0)
71
-
72
-
73
-@pytest.mark.skipif(torch.cuda.device_count() < 2,
74
- reason="Need at least 2 GPUs to run the test.")
75
-@pytest.mark.parametrize(
76
- "common_llm_kwargs",
77
- [[
78
- # Skip cuda graph recording for fast test.
79
- "--enforce-eager",
80
- "--tensor_parallel_size",
81
- "2",
82
-
83
- # precision
84
- "--dtype",
85
- "bfloat16",
86
- ]])
87
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [[]])
88
-@pytest.mark.parametrize("baseline_llm_kwargs", [[]])
89
-@pytest.mark.parametrize(
90
- "model, test_llm_kwargs",
91
- [("JackFram/llama-68m", [
92
- "--speculative_config",
93
- json.dumps({
94
- "model": "JackFram/llama-68m",
95
- "num_speculative_tokens": 5,
96
- "draft_tensor_parallel_size": 1,
97
- }),
98
- ]),
99
- ("ibm-granite/granite-3b-code-instruct", [
100
- "--speculative_config",
101
- json.dumps({
102
- "model": "ibm-granite/granite-3b-code-instruct",
103
- "num_speculative_tokens": 5,
104
- "draft_tensor_parallel_size": 1,
105
- }),
106
- ])])
107
-@pytest.mark.parametrize("batch_size", [2])
108
-@pytest.mark.parametrize("seed", [1])
109
-def test_draft_model_tp_lt_target_model_tp2(model, common_llm_kwargs,
110
- per_test_common_llm_kwargs,
111
- baseline_llm_kwargs,
112
- test_llm_kwargs, batch_size: int,
113
- seed: int):
114
- """Verify spec decode works well with smaller tp for draft models.
115
- """
116
- run_equality_correctness_test_tp(model,
117
- common_llm_kwargs,
118
- per_test_common_llm_kwargs,
119
- baseline_llm_kwargs,
120
- test_llm_kwargs,
121
- batch_size,
122
- max_output_len=32,
123
- seed=seed,
124
- temperature=0.0)
125
-
126
-
127
-@pytest.mark.skipif(torch.cuda.device_count() < 2,
128
- reason="Need at least 2 GPUs to run the test.")
129
-@pytest.mark.parametrize(
130
- "common_llm_kwargs",
131
- [[
132
- # Skip cuda graph recording for fast test.
133
- "--enforce-eager",
134
- "--tensor_parallel_size",
135
- "2",
136
-
137
- # precision
138
- "--dtype",
139
- "bfloat16",
140
- ]])
141
-@pytest.mark.parametrize(
142
- "per_test_common_llm_kwargs",
143
- [["--enable-chunked-prefill", "False"],
144
- [
145
- "--enable-chunked-prefill", "True", "--max-num-batched-tokens", "4",
146
- "--max-num-seqs", "4"
147
- ]])
148
-@pytest.mark.parametrize("baseline_llm_kwargs", [[]])
149
-@pytest.mark.parametrize("model, test_llm_kwargs",
150
- [("JackFram/llama-68m", [
151
- "--speculative_config",
152
- json.dumps({
153
- "model": "JackFram/llama-68m",
154
- "num_speculative_tokens": 3,
155
- }),
156
- ]),
157
- ("JackFram/llama-68m", [
158
- "--speculative_config",
159
- json.dumps({
160
- "model": "JackFram/llama-68m",
161
- "num_speculative_tokens": 3,
162
- "draft_tensor_parallel_size": 1,
163
- }),
164
- ])])
165
-@pytest.mark.parametrize("logprobs", [None])
166
-@pytest.mark.parametrize("batch_size", [2])
167
-@pytest.mark.parametrize("seed", [1])
168
-def test_spec_decode_chunked_prefill_tp2(model, common_llm_kwargs,
169
- per_test_common_llm_kwargs,
170
- baseline_llm_kwargs, test_llm_kwargs,
171
- logprobs: Optional[int],
172
- batch_size: int, seed: int):
173
- """Verify spec decode works well with same and different TP size for
174
- the draft model with chunked prefill.
175
- """
176
- run_equality_correctness_test_tp(model,
177
- common_llm_kwargs,
178
- per_test_common_llm_kwargs,
179
- baseline_llm_kwargs,
180
- test_llm_kwargs,
181
- batch_size,
182
- max_output_len=32,
183
- seed=seed,
184
- temperature=0.0,
185
- logprobs=logprobs)
186
-
187
-
188
-@pytest.mark.skipif(torch.cuda.device_count() < 2,
189
- reason="Need at least 2 GPUs to run the test.")
190
-@pytest.mark.parametrize(
191
- "common_llm_kwargs",
192
- [[
193
- # Skip cuda graph recording for fast test.
194
- "--enforce-eager",
195
- "--tensor_parallel_size",
196
- "2",
197
-
198
- # precision
199
- "--dtype",
200
- "bfloat16",
201
- ]])
202
-@pytest.mark.parametrize(
203
- "per_test_common_llm_kwargs",
204
- [["--enable-chunked-prefill", "False"],
205
- [
206
- "--enable-chunked-prefill", "True", "--max-num-batched-tokens", "4",
207
- "--max-num-seqs", "4"
208
- ]])
209
-@pytest.mark.parametrize("baseline_llm_kwargs", [[]])
210
-@pytest.mark.parametrize("model, test_llm_kwargs",
211
- [("JackFram/llama-68m", [
212
- "--speculative_config",
213
- json.dumps({
214
- "model": "JackFram/llama-68m",
215
- "num_speculative_tokens": 3,
216
- "disable_logprobs": False,
217
- }),
218
- ]),
219
- ("JackFram/llama-68m", [
220
- "--speculative_config",
221
- json.dumps({
222
- "model": "JackFram/llama-68m",
223
- "num_speculative_tokens": 3,
224
- "draft_tensor_parallel_size": 1,
225
- "disable_logprobs": False,
226
- }),
227
- ])])
228
-@pytest.mark.parametrize("logprobs", [2])
229
-@pytest.mark.parametrize("batch_size", [2])
230
-@pytest.mark.parametrize("seed", [1])
231
-def test_spec_decode_chunked_prefill_tp2_with_logprobs(
232
- model, common_llm_kwargs, per_test_common_llm_kwargs,
233
- baseline_llm_kwargs, test_llm_kwargs, logprobs: Optional[int],
234
- batch_size: int, seed: int):
235
- """Verify spec decode works well with same and different TP size for
236
- the draft model with chunked prefill.
237
- """
238
- run_equality_correctness_test_tp(model,
239
- common_llm_kwargs,
240
- per_test_common_llm_kwargs,
241
- baseline_llm_kwargs,
242
- test_llm_kwargs,
243
- batch_size,
244
- max_output_len=32,
245
- seed=seed,
246
- temperature=0.0,
247
- logprobs=logprobs)
tests/spec_decode/e2e/test_integration_dist_tp4.py
deleted
-123
@@ -1,123 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests which cover integration of the speculative decoding framework with
4
-tensor parallelism.
5
-"""
6
-
7
-import json
8
-
9
-import openai
10
-import pytest
11
-import torch
12
-
13
-from .conftest import run_equality_correctness_test_tp
14
-
15
-MAIN_MODEL = "JackFram/llama-68m"
16
-SPEC_MODEL = "JackFram/llama-68m"
17
-
18
-
19
-@pytest.mark.skipif(torch.cuda.device_count() < 4,
20
- reason="Need at least 4 GPUs to run the test.")
21
-@pytest.mark.parametrize(
22
- "common_llm_kwargs",
23
- [[
24
- # Skip cuda graph recording for fast test.
25
- "--enforce_eager",
26
- "--tensor-parallel-size",
27
- "4",
28
- ]])
29
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [
30
- [],
31
-])
32
-@pytest.mark.parametrize("baseline_llm_kwargs", [[]])
33
-@pytest.mark.parametrize(
34
- "test_llm_kwargs",
35
- [
36
- #TODO(wooyeon): add spec_draft_dp=2 case
37
- [
38
- "--speculative_config",
39
- json.dumps({
40
- "model": f"{SPEC_MODEL}",
41
- "num_speculative_tokens": 5,
42
- "draft_tensor_parallel_size": 1,
43
- }),
44
- ],
45
- ])
46
-@pytest.mark.parametrize("batch_size", [2])
47
-@pytest.mark.parametrize("seed", [1])
48
-def test_draft_model_tp_lt_target_model_tp4(common_llm_kwargs,
49
- per_test_common_llm_kwargs,
50
- baseline_llm_kwargs,
51
- test_llm_kwargs, batch_size: int,
52
- seed: int):
53
- """Verify spec decode works well with smaller tp for draft models.
54
- """
55
- run_equality_correctness_test_tp(MAIN_MODEL,
56
- common_llm_kwargs,
57
- per_test_common_llm_kwargs,
58
- baseline_llm_kwargs,
59
- test_llm_kwargs,
60
- batch_size,
61
- max_output_len=32,
62
- seed=seed,
63
- temperature=0.0)
64
-
65
-
66
-@pytest.mark.skipif(torch.cuda.device_count() < 4,
67
- reason="Need at least 4 GPUs to run the test.")
68
-@pytest.mark.parametrize(
69
- "common_llm_kwargs",
70
- [[
71
-
72
- # Skip cuda graph recording for fast test.
73
- "--enforce-eager",
74
- "--tensor-parallel-size",
75
- "4",
76
- ]])
77
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [[]])
78
-@pytest.mark.parametrize("baseline_llm_kwargs", [[]])
79
-@pytest.mark.parametrize(
80
- "test_llm_kwargs",
81
- [
82
- [
83
- # Artificially limit the draft model max model len; this forces vLLM
84
- # to skip speculation once the sequences grow beyond 32-k tokens.
85
- "--speculative_config",
86
- json.dumps({
87
- "model": f"{SPEC_MODEL}",
88
- "num_speculative_tokens": 5,
89
- "max_model_len": 32,
90
- }),
91
- ],
92
- ])
93
-@pytest.mark.parametrize("batch_size", [8])
94
-@pytest.mark.parametrize(
95
- "output_len",
96
- [
97
- # This must be a good bit larger than speculative_max_model_len so that
98
- # we can test the case where all seqs are skipped, but still small to
99
- # ensure fast test.
100
- 64,
101
- ])
102
-@pytest.mark.parametrize("seed", [1])
103
-def test_skip_speculation(common_llm_kwargs, per_test_common_llm_kwargs,
104
- baseline_llm_kwargs, test_llm_kwargs,
105
- batch_size: int, output_len: int, seed: int):
106
- """Verify job failure with RuntimeError when all sequences skip speculation.
107
- We do this by setting the max model len of the draft model to an
108
- artificially low value, such that when the sequences grow beyond it, they
109
- are skipped in speculative decoding.
110
-
111
- TODO: fix it to pass without raising Error. (#5814)
112
- """
113
- with pytest.raises(
114
- (openai.APIConnectionError, openai.InternalServerError)):
115
- run_equality_correctness_test_tp(MAIN_MODEL,
116
- common_llm_kwargs,
117
- per_test_common_llm_kwargs,
118
- baseline_llm_kwargs,
119
- test_llm_kwargs,
120
- batch_size,
121
- output_len,
122
- seed,
123
- temperature=0.0)
tests/spec_decode/e2e/test_logprobs.py
deleted
-315
@@ -1,315 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from itertools import cycle
5
-
6
-import pytest
7
-
8
-from vllm import SamplingParams
9
-
10
-from ..utils import maybe_enable_chunked_prefill
11
-from .conftest import run_equality_correctness_test
12
-
13
-
14
-@pytest.mark.parametrize(
15
- "common_llm_kwargs",
16
- [{
17
- "model_name": "JackFram/llama-160m",
18
-
19
- # Skip cuda graph recording for fast test.
20
- "enforce_eager": True,
21
-
22
- # The original model is float32, keep it for numerical stability.
23
- "dtype": "float32",
24
- }])
25
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
26
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
27
-@pytest.mark.parametrize("test_llm_kwargs", [{
28
- "speculative_config": {
29
- "model": "JackFram/llama-68m",
30
- "num_speculative_tokens": 3,
31
- "disable_logprobs": False,
32
- },
33
-}, {
34
- "speculative_config": {
35
- "model": "JackFram/llama-68m",
36
- "num_speculative_tokens": 3,
37
- "disable_logprobs": True,
38
- },
39
-}])
40
-@pytest.mark.parametrize("batch_size", [8])
41
-@pytest.mark.parametrize(
42
- "output_len",
43
- [
44
- # Use smaller output len for fast test.
45
- 7,
46
- ])
47
-@pytest.mark.parametrize("seed", [1])
48
-@pytest.mark.parametrize("logprobs", [1, 6])
49
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 4, 12])
50
-def test_logprobs_equality(vllm_runner, common_llm_kwargs,
51
- per_test_common_llm_kwargs, baseline_llm_kwargs,
52
- test_llm_kwargs, batch_size: int, output_len: int,
53
- seed: int, logprobs: int, prefill_chunk_size: int):
54
- """Verify output logprobs are equal with and without speculative decoding,
55
- as well as with and without chunked prefill.
56
- """
57
- maybe_enable_chunked_prefill(prefill_chunk_size, common_llm_kwargs)
58
- run_equality_correctness_test(
59
- vllm_runner,
60
- common_llm_kwargs,
61
- per_test_common_llm_kwargs,
62
- baseline_llm_kwargs,
63
- test_llm_kwargs,
64
- batch_size,
65
- output_len,
66
- seed,
67
- temperature=0.0,
68
- logprobs=logprobs,
69
- prompt_logprobs=logprobs,
70
- disable_logprobs=test_llm_kwargs["speculative_config"]
71
- ["disable_logprobs"])
72
-
73
-
74
-@pytest.mark.parametrize(
75
- "common_llm_kwargs",
76
- [{
77
- "model_name": "JackFram/llama-68m",
78
-
79
- # Skip cuda graph recording for fast test.
80
- "enforce_eager": True,
81
-
82
- # The original model is float32, keep it for numerical stability.
83
- "dtype": "float32",
84
- }])
85
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
86
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
87
-@pytest.mark.parametrize("test_llm_kwargs", [{
88
- "speculative_config": {
89
- "model": "JackFram/llama-160m",
90
- "num_speculative_tokens": 3,
91
- "disable_logprobs": False,
92
- },
93
-}, {
94
- "speculative_config": {
95
- "model": "JackFram/llama-160m",
96
- "num_speculative_tokens": 6,
97
- "disable_logprobs": False,
98
- },
99
-}])
100
-@pytest.mark.parametrize("batch_size", [8])
101
-@pytest.mark.parametrize(
102
- "output_len",
103
- [
104
- # Use smaller output len for fast test.
105
- 32,
106
- ])
107
-@pytest.mark.parametrize("seed", [1])
108
-@pytest.mark.parametrize("logprobs", [1, 6])
109
-def test_logprobs_different_k(vllm_runner, common_llm_kwargs,
110
- per_test_common_llm_kwargs, baseline_llm_kwargs,
111
- test_llm_kwargs, batch_size: int,
112
- output_len: int, seed: int, logprobs: int):
113
- """Veriy logprob greedy equality with different speculation lens.
114
- """
115
- run_equality_correctness_test(
116
- vllm_runner,
117
- common_llm_kwargs,
118
- per_test_common_llm_kwargs,
119
- baseline_llm_kwargs,
120
- test_llm_kwargs,
121
- batch_size,
122
- output_len,
123
- seed,
124
- temperature=0.0,
125
- logprobs=logprobs,
126
- disable_logprobs=test_llm_kwargs["speculative_config"]
127
- ["disable_logprobs"])
128
-
129
-
130
-@pytest.mark.parametrize(
131
- "common_llm_kwargs",
132
- [{
133
- "model_name": "JackFram/llama-68m",
134
-
135
- # Skip cuda graph recording for fast test.
136
- "enforce_eager": True,
137
-
138
- # The original model is float32, keep it for numerical stability.
139
- "dtype": "float32",
140
- }])
141
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
142
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
143
-@pytest.mark.parametrize(
144
- "test_llm_kwargs",
145
- [{
146
- "speculative_config": {
147
- "model": "JackFram/llama-160m",
148
- "num_speculative_tokens": 3,
149
- "disable_logprobs": False,
150
- # Artificially limit the draft model max model len; this forces
151
- # vLLM to skip speculation once the sequences grow beyond 32-k
152
- # tokens.
153
- "max_model_len": 32,
154
- },
155
- }])
156
-@pytest.mark.parametrize("batch_size", [8])
157
-@pytest.mark.parametrize(
158
- "output_len",
159
- [
160
- # Use smaller output len for fast test.
161
- 32,
162
- ])
163
-@pytest.mark.parametrize("seed", [1])
164
-@pytest.mark.parametrize("logprobs", [1])
165
-def test_logprobs_when_skip_speculation(vllm_runner, common_llm_kwargs,
166
- per_test_common_llm_kwargs,
167
- baseline_llm_kwargs, test_llm_kwargs,
168
- batch_size: int, output_len: int,
169
- seed: int, logprobs: int):
170
- """Verify logprobs greedy equality when some sequences skip speculation.
171
- """
172
- run_equality_correctness_test(
173
- vllm_runner,
174
- common_llm_kwargs,
175
- per_test_common_llm_kwargs,
176
- baseline_llm_kwargs,
177
- test_llm_kwargs,
178
- batch_size,
179
- output_len,
180
- seed,
181
- temperature=0.0,
182
- logprobs=logprobs,
183
- disable_logprobs=test_llm_kwargs["speculative_config"]
184
- ["disable_logprobs"])
185
-
186
-
187
-@pytest.mark.parametrize(
188
- "common_llm_kwargs",
189
- [{
190
- "model_name": "JackFram/llama-68m",
191
-
192
- # Skip cuda graph recording for fast test.
193
- "enforce_eager": True,
194
-
195
- # The original model is float32, keep it for numerical stability.
196
- "dtype": "float32",
197
- }])
198
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
199
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
200
-@pytest.mark.parametrize("test_llm_kwargs", [{
201
- "speculative_config": {
202
- "model": "JackFram/llama-160m",
203
- "num_speculative_tokens": 3,
204
- "disable_logprobs": False,
205
- },
206
-}])
207
-@pytest.mark.parametrize("batch_size", [1])
208
-@pytest.mark.parametrize(
209
- "output_len",
210
- [
211
- # Use smaller output len for fast test.
212
- 32,
213
- ])
214
-@pytest.mark.parametrize("seed", [1])
215
-@pytest.mark.parametrize("logprobs", [6])
216
-def test_logprobs_temp_1(vllm_runner, common_llm_kwargs,
217
- per_test_common_llm_kwargs, baseline_llm_kwargs,
218
- test_llm_kwargs, batch_size: int, output_len: int,
219
- seed: int, logprobs: int):
220
- """Verify at least one logprob result has num_logprobs+1, which tests the
221
- case where the sampled token is not in top-k logprobs.
222
-
223
- Ideally, this test should validate equality with non-spec by getting
224
- logprobs. This is left as future improvement.
225
- """
226
- temperature = 1.0
227
-
228
- prompts = [
229
- "Hello, my name is",
230
- "The president of the United States is",
231
- "The capital of France is",
232
- "The future of AI is",
233
- "San Francisco is know for its",
234
- "Facebook was created in 2004 by",
235
- "Curious George is a",
236
- "Python 3.11 brings improvements to its",
237
- ]
238
-
239
- prompts = [prompt for prompt, _ in zip(cycle(prompts), range(batch_size))]
240
-
241
- sampling_params = SamplingParams(
242
- max_tokens=output_len,
243
- ignore_eos=True,
244
- temperature=temperature,
245
- logprobs=logprobs,
246
- )
247
-
248
- sd_args = {
249
- **common_llm_kwargs,
250
- **per_test_common_llm_kwargs,
251
- **test_llm_kwargs,
252
- }
253
-
254
- with vllm_runner(**sd_args) as vllm_model:
255
- sd_outputs = vllm_model.generate_w_logprobs(prompts, sampling_params)
256
-
257
- num_returned_logprobs = [
258
- len(seq_logprobs) for seq_logprobs in sd_outputs[-1]
259
- ]
260
-
261
- # Assert one of the returned logprobs has > num_logprobs (indicating the
262
- # sampled token is not in top-k).
263
- assert any(
264
- [num_returned > logprobs for num_returned in num_returned_logprobs])
265
-
266
-
267
-@pytest.mark.parametrize(
268
- "common_llm_kwargs",
269
- [{
270
- "model_name": "JackFram/llama-160m",
271
-
272
- # Skip cuda graph recording for fast test.
273
- "enforce_eager": True,
274
-
275
- # The original model is float32, keep it for numerical stability.
276
- "dtype": "float32",
277
- }])
278
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
279
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
280
-@pytest.mark.parametrize("test_llm_kwargs", [{
281
- "speculative_config": {
282
- "model": "JackFram/llama-68m",
283
- "num_speculative_tokens": 3,
284
- "disable_logprobs": True,
285
- },
286
-}])
287
-@pytest.mark.parametrize("seed", [1])
288
-@pytest.mark.parametrize("batch_size", [4])
289
-@pytest.mark.parametrize(
290
- "output_len",
291
- [
292
- # Use smaller output len for fast test.
293
- 32,
294
- ])
295
-@pytest.mark.parametrize("logprobs", [0])
296
-def test_logprobs_disabled(vllm_runner, common_llm_kwargs,
297
- per_test_common_llm_kwargs, baseline_llm_kwargs,
298
- test_llm_kwargs, batch_size: int, output_len: int,
299
- seed: int, logprobs: int):
300
- """Check the behavior when logprobs are disabled.
301
- Token choices should match with the base model.
302
- """
303
- run_equality_correctness_test(
304
- vllm_runner,
305
- common_llm_kwargs,
306
- per_test_common_llm_kwargs,
307
- baseline_llm_kwargs,
308
- test_llm_kwargs,
309
- batch_size,
310
- output_len,
311
- seed,
312
- temperature=0.0,
313
- logprobs=logprobs,
314
- disable_logprobs=test_llm_kwargs["speculative_config"]
315
- ["disable_logprobs"])
tests/spec_decode/e2e/test_medusa_correctness.py
deleted
-417
@@ -1,417 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""This docstring details important information on the testing methodology.
4
-
5
-Most of the tests rely on "greedy equality", where we expect the output of
6
-speculative decoding on a sequence to exactly match the output of normal non-
7
-speculative decoding.
8
-
9
-Since speculative decoding with rejection sampling guarantees that the output
10
-distribution matches the target model's output distribution (up to hardware
11
-numerics, see https://arxiv.org/pdf/2302.01318.pdf), we can expect greedy
12
-equality.
13
-
14
-However, we still need to verify below scenario could be passed:
15
- * Batch size 1 greedy equality
16
- * Batch size >1 greedy equality
17
- * Test greedy equality under preemption
18
- * Test greedy equality under various number of speculative tokens.
19
-
20
-With those tests, we can say at least, Medusa would not break the
21
-correctness for the target model outputs.
22
-"""
23
-
24
-import pytest
25
-
26
-from ..utils import maybe_enable_chunked_prefill
27
-from .conftest import run_equality_correctness_test
28
-
29
-# main model
30
-# lmsys/vicuna-7b-v1.3 was to be used but it's causing
31
-# OOM in CI pipeline, so using a smaller model.
32
-MAIN_MODEL = "JackFram/llama-68m"
33
-
34
-# speculative model
35
-SPEC_MODEL = "abhigoyal/vllm-medusa-llama-68m-random"
36
-
37
-# max number of speculative tokens: this corresponds to
38
-# num_heads in the config.json of the speculator model.
39
-MAX_SPEC_TOKENS = 5
40
-
41
-# precision
42
-PRECISION = "float32"
43
-
44
-
45
-@pytest.mark.parametrize(
46
- "common_llm_kwargs",
47
- [{
48
- # Skip cuda graph recording for fast test.
49
- "enforce_eager": True,
50
-
51
- # Print spec metrics.
52
- "disable_log_stats": False,
53
-
54
- # Precision
55
- "dtype": PRECISION,
56
-
57
- # Main model
58
- "model_name": MAIN_MODEL,
59
- }])
60
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
61
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
62
-@pytest.mark.parametrize("test_llm_kwargs", [
63
- {
64
- "speculative_config": {
65
- "model": SPEC_MODEL,
66
- "num_speculative_tokens": MAX_SPEC_TOKENS,
67
- },
68
- },
69
-])
70
-@pytest.mark.parametrize("output_len", [
71
- 128,
72
-])
73
-@pytest.mark.parametrize("batch_size", [1, 32])
74
-@pytest.mark.parametrize("seed", [1])
75
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 32])
76
-def test_medusa_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
77
- per_test_common_llm_kwargs,
78
- baseline_llm_kwargs, test_llm_kwargs,
79
- batch_size: int, output_len: int,
80
- seed: int, prefill_chunk_size: int):
81
- """Verify greedy equality with different batch size."""
82
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
83
- run_equality_correctness_test(vllm_runner,
84
- common_llm_kwargs,
85
- per_test_common_llm_kwargs,
86
- baseline_llm_kwargs,
87
- test_llm_kwargs,
88
- batch_size,
89
- max_output_len=output_len,
90
- seed=seed,
91
- temperature=0.0)
92
-
93
-
94
-@pytest.mark.parametrize(
95
- "common_llm_kwargs",
96
- [{
97
- # Skip cuda graph recording for fast test.
98
- "enforce_eager": True,
99
-
100
- # Print spec metrics.
101
- "disable_log_stats": False,
102
-
103
- # Precision
104
- "dtype": PRECISION,
105
-
106
- # Main model
107
- "model_name": MAIN_MODEL,
108
- }])
109
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
110
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
111
-@pytest.mark.parametrize("test_llm_kwargs", [
112
- {
113
- "speculative_config": {
114
- "model": SPEC_MODEL,
115
- "num_speculative_tokens": MAX_SPEC_TOKENS,
116
- "disable_logprobs": False,
117
- },
118
- },
119
- {
120
- "speculative_config": {
121
- "model": SPEC_MODEL,
122
- "num_speculative_tokens": MAX_SPEC_TOKENS,
123
- "disable_logprobs": True,
124
- },
125
- },
126
-])
127
-@pytest.mark.parametrize("output_len", [
128
- 8,
129
-])
130
-@pytest.mark.parametrize("batch_size", [8])
131
-@pytest.mark.parametrize("seed", [1])
132
-@pytest.mark.parametrize("logprobs", [1, 6])
133
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 32])
134
-def test_medusa_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
135
- per_test_common_llm_kwargs,
136
- baseline_llm_kwargs, test_llm_kwargs,
137
- batch_size: int, output_len: int,
138
- seed: int, logprobs: int,
139
- prefill_chunk_size: int):
140
- """Verify greedy equality with different batch size."""
141
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
142
- run_equality_correctness_test(
143
- vllm_runner,
144
- common_llm_kwargs,
145
- per_test_common_llm_kwargs,
146
- baseline_llm_kwargs,
147
- test_llm_kwargs,
148
- batch_size,
149
- max_output_len=output_len,
150
- seed=seed,
151
- temperature=0.0,
152
- logprobs=logprobs,
153
- prompt_logprobs=logprobs,
154
- disable_logprobs=test_llm_kwargs["speculative_config"]
155
- ["disable_logprobs"])
156
-
157
-
158
-@pytest.mark.parametrize(
159
- "common_llm_kwargs",
160
- [{
161
- "enforce_eager": False,
162
-
163
- # Print spec metrics.
164
- "disable_log_stats": False,
165
-
166
- # Precision
167
- "dtype": PRECISION,
168
-
169
- # Main model
170
- "model_name": MAIN_MODEL,
171
- }])
172
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
173
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
174
-@pytest.mark.parametrize("test_llm_kwargs", [
175
- {
176
- "speculative_config": {
177
- "model": SPEC_MODEL,
178
- "num_speculative_tokens": MAX_SPEC_TOKENS,
179
- },
180
- },
181
-])
182
-@pytest.mark.parametrize("output_len", [
183
- 128,
184
-])
185
-@pytest.mark.parametrize("batch_size", [1, 32])
186
-@pytest.mark.parametrize("seed", [1])
187
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 32])
188
-def test_medusa_e2e_greedy_correctness_cuda_graph(
189
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
190
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
191
- seed: int, prefill_chunk_size: int):
192
- """Verify greedy equality with cuda graph enabled and different
193
- batch sizes."""
194
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
195
- run_equality_correctness_test(vllm_runner,
196
- common_llm_kwargs,
197
- per_test_common_llm_kwargs,
198
- baseline_llm_kwargs,
199
- test_llm_kwargs,
200
- batch_size,
201
- max_output_len=output_len,
202
- seed=seed,
203
- temperature=0.0)
204
-
205
-
206
-@pytest.mark.parametrize(
207
- "common_llm_kwargs",
208
- [{
209
- "block_size": 16,
210
- # 2 for small prompt, 256//8 for generated.
211
- "num_gpu_blocks_override": 2 + 256 // 8,
212
- "max_model_len": (2 + 256 // 8) * 8,
213
-
214
- # Skip cuda graph recording for fast test.
215
- "enforce_eager": True,
216
-
217
- # Precision
218
- "dtype": PRECISION,
219
-
220
- # Main model
221
- "model_name": MAIN_MODEL,
222
- }])
223
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
224
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
225
-@pytest.mark.parametrize("test_llm_kwargs", [
226
- {
227
- "speculative_config": {
228
- "model": SPEC_MODEL,
229
- "num_speculative_tokens": MAX_SPEC_TOKENS,
230
- },
231
- },
232
-])
233
-@pytest.mark.parametrize(
234
- "output_len",
235
- [
236
- # Use small output len for fast test.
237
- 128,
238
- ])
239
-@pytest.mark.parametrize("batch_size", [4])
240
-@pytest.mark.parametrize("seed", [1])
241
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 32])
242
-def test_medusa_e2e_greedy_correctness_with_preemption(
243
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
244
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
245
- seed: int, prefill_chunk_size: int):
246
- """Verify greedy equality, even when some sequences are preempted mid-
247
- generation.
248
- """
249
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
250
- run_equality_correctness_test(vllm_runner,
251
- common_llm_kwargs,
252
- per_test_common_llm_kwargs,
253
- baseline_llm_kwargs,
254
- test_llm_kwargs,
255
- batch_size,
256
- max_output_len=output_len,
257
- seed=seed,
258
- temperature=0.0)
259
-
260
-
261
-@pytest.mark.parametrize(
262
- "common_llm_kwargs",
263
- [{
264
- # Skip cuda graph recording for fast test.
265
- "enforce_eager": True,
266
-
267
- # Precision
268
- "dtype": PRECISION,
269
-
270
- # Main model
271
- "model_name": MAIN_MODEL,
272
- }])
273
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
274
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
275
-@pytest.mark.parametrize(
276
- "test_llm_kwargs",
277
- [
278
- {
279
- "speculative_config": {
280
- "model": SPEC_MODEL,
281
- "num_speculative_tokens": k,
282
- },
283
- }
284
- # Try a range of num. speculative tokens
285
- for k in range(1, 1 + MAX_SPEC_TOKENS)
286
- ])
287
-@pytest.mark.parametrize("batch_size", [2])
288
-@pytest.mark.parametrize(
289
- "output_len",
290
- [
291
- # Use smaller output len for fast test.
292
- 32,
293
- ])
294
-@pytest.mark.parametrize("seed", [1])
295
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 32])
296
-def test_medusa_different_k(vllm_runner, common_llm_kwargs,
297
- per_test_common_llm_kwargs, baseline_llm_kwargs,
298
- test_llm_kwargs, batch_size: int, output_len: int,
299
- seed: int, prefill_chunk_size: int):
300
- """Verify that medusa speculative decoding produces exact equality
301
- to without spec decode with different values of num_speculative_tokens.
302
- """
303
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
304
- run_equality_correctness_test(vllm_runner,
305
- common_llm_kwargs,
306
- per_test_common_llm_kwargs,
307
- baseline_llm_kwargs,
308
- test_llm_kwargs,
309
- batch_size,
310
- max_output_len=output_len,
311
- seed=seed,
312
- temperature=0.0)
313
-
314
-
315
-@pytest.mark.parametrize(
316
- "common_llm_kwargs",
317
- [{
318
- # Skip cuda graph recording for fast test.
319
- "enforce_eager": True,
320
-
321
- # Precision
322
- "dtype": PRECISION,
323
-
324
- # Main model
325
- "model_name": MAIN_MODEL,
326
- }])
327
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
328
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
329
-@pytest.mark.parametrize("test_llm_kwargs", [{
330
- "speculative_config": {
331
- "model": SPEC_MODEL,
332
- "num_speculative_tokens": MAX_SPEC_TOKENS,
333
- "disable_by_batch_size": 4,
334
- },
335
-}])
336
-@pytest.mark.parametrize("batch_size", [1, 5])
337
-@pytest.mark.parametrize(
338
- "output_len",
339
- [
340
- # Use smaller output len for fast test.
341
- 32,
342
- ])
343
-@pytest.mark.parametrize("seed", [1])
344
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 32])
345
-def test_medusa_disable_queue(vllm_runner, common_llm_kwargs,
346
- per_test_common_llm_kwargs, baseline_llm_kwargs,
347
- test_llm_kwargs, batch_size: int,
348
- output_len: int, seed: int,
349
- prefill_chunk_size: int):
350
- """Verify that medusa speculative decoding produces exact equality
351
- to without spec decode when speculation is disabled for large
352
- batch sizes.
353
- """
354
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
355
- run_equality_correctness_test(vllm_runner,
356
- common_llm_kwargs,
357
- per_test_common_llm_kwargs,
358
- baseline_llm_kwargs,
359
- test_llm_kwargs,
360
- batch_size,
361
- max_output_len=output_len,
362
- seed=seed,
363
- temperature=0.0)
364
-
365
-
366
-@pytest.mark.parametrize(
367
- "common_llm_kwargs",
368
- [{
369
- # Skip cuda graph recording for fast test.
370
- "enforce_eager": True,
371
-
372
- # Precision
373
- "dtype": PRECISION,
374
-
375
- # Main model
376
- "model_name": MAIN_MODEL,
377
- }])
378
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
379
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
380
-@pytest.mark.parametrize("test_llm_kwargs", [{
381
- "speculative_config": {
382
- "model": SPEC_MODEL,
383
- "num_speculative_tokens": MAX_SPEC_TOKENS,
384
- "disable_by_batch_size": 4,
385
- "disable_mqa_scorer": True,
386
- },
387
-}])
388
-@pytest.mark.parametrize("batch_size", [1, 5])
389
-@pytest.mark.parametrize(
390
- "output_len",
391
- [
392
- # Use smaller output len for fast test.
393
- 32,
394
- ])
395
-@pytest.mark.parametrize("seed", [1])
396
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 32])
397
-def test_mqa_scorer(vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
398
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int,
399
- output_len: int, seed: int, prefill_chunk_size: int):
400
- """Verify that speculative decoding generates the same output
401
- with batch expansion scorer and mqa scorer.
402
- """
403
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
404
- run_equality_correctness_test(vllm_runner,
405
- common_llm_kwargs,
406
- per_test_common_llm_kwargs,
407
- baseline_llm_kwargs,
408
- test_llm_kwargs,
409
- batch_size,
410
- max_output_len=output_len,
411
- seed=seed,
412
- temperature=0.0)
413
-
414
-
415
-if __name__ == "__main__":
416
- import pytest
417
- pytest.main([__file__])
tests/spec_decode/e2e/test_mlp_correctness.py
deleted
-533
@@ -1,533 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""This docstring details important information on the testing methodology.
4
-
5
-Most of the tests rely on "greedy equality", where we expect the output of
6
-speculative decoding on a sequence to exactly match the output of normal non-
7
-speculative decoding.
8
-
9
-Since speculative decoding with rejection sampling guarantees that the output
10
-distribution matches the target model's output distribution (up to hardware
11
-numerics, see https://arxiv.org/pdf/2302.01318.pdf), we can expect greedy
12
-equality.
13
-
14
-However, we still need to verify below scenario could be passed:
15
- * Batch size 1 greedy equality
16
- * Batch size >1 greedy equality
17
- * Test greedy equality under preemption
18
- * Test greedy equality under various number of speculative tokens.
19
-
20
-With those tests, we can say at least, MLPSpeculator would not break the
21
-correctness for the target model outputs.
22
-"""
23
-
24
-from unittest.mock import patch
25
-
26
-import pytest
27
-
28
-from vllm.model_executor.layers.vocab_parallel_embedding import pad_vocab_size
29
-
30
-from ..utils import maybe_enable_chunked_prefill
31
-from .conftest import run_equality_correctness_test
32
-
33
-# main model
34
-MAIN_MODEL = "JackFram/llama-160m"
35
-
36
-# speculative model
37
-SPEC_MODEL = "ibm-ai-platform/llama-160m-accelerator"
38
-
39
-# max. number of speculative tokens: this corresponds to
40
-# n_predict in the config.json of the speculator model.
41
-MAX_SPEC_TOKENS = 3
42
-
43
-# precision
44
-PRECISION = "float32"
45
-
46
-
47
-@pytest.mark.parametrize(
48
- "common_llm_kwargs",
49
- [{
50
- # Skip cuda graph recording for fast test.
51
- "enforce_eager": True,
52
-
53
- # Print spec metrics.
54
- "disable_log_stats": False,
55
-
56
- # Precision
57
- "dtype": PRECISION,
58
-
59
- # Main model
60
- "model_name": MAIN_MODEL,
61
- }])
62
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
63
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
64
-@pytest.mark.parametrize("test_llm_kwargs", [
65
- {
66
- "speculative_config": {
67
- "model": SPEC_MODEL,
68
- },
69
- },
70
-])
71
-@pytest.mark.parametrize("output_len", [
72
- 128,
73
-])
74
-@pytest.mark.parametrize("batch_size", [4, 32])
75
-@pytest.mark.parametrize("seed", [1])
76
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 32])
77
-def test_mlp_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
78
- per_test_common_llm_kwargs,
79
- baseline_llm_kwargs, test_llm_kwargs,
80
- batch_size: int, output_len: int,
81
- seed: int, prefill_chunk_size: int):
82
- """Verify greedy equality with different batch size."""
83
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
84
- run_equality_correctness_test(vllm_runner,
85
- common_llm_kwargs,
86
- per_test_common_llm_kwargs,
87
- baseline_llm_kwargs,
88
- test_llm_kwargs,
89
- batch_size,
90
- max_output_len=output_len,
91
- seed=seed,
92
- temperature=0.0)
93
-
94
-
95
-@pytest.mark.parametrize(
96
- "common_llm_kwargs",
97
- [{
98
- # Skip cuda graph recording for fast test.
99
- "enforce_eager": True,
100
-
101
- # Print spec metrics.
102
- "disable_log_stats": False,
103
-
104
- # Precision
105
- "dtype": PRECISION,
106
-
107
- # Main model
108
- "model_name": MAIN_MODEL,
109
- }])
110
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
111
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
112
-@pytest.mark.parametrize("test_llm_kwargs", [
113
- {
114
- "speculative_config": {
115
- "model": SPEC_MODEL,
116
- "disable_logprobs": False,
117
- },
118
- },
119
- {
120
- "speculative_config": {
121
- "model": SPEC_MODEL,
122
- "disable_logprobs": True,
123
- },
124
- },
125
-])
126
-@pytest.mark.parametrize("output_len", [8])
127
-@pytest.mark.parametrize("batch_size", [8])
128
-@pytest.mark.parametrize("seed", [1])
129
-@pytest.mark.parametrize("logprobs", [1, 6])
130
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 4])
131
-def test_mlp_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
132
- per_test_common_llm_kwargs,
133
- baseline_llm_kwargs, test_llm_kwargs,
134
- batch_size: int, output_len: int, seed: int,
135
- logprobs: int, prefill_chunk_size: int):
136
- """Verify greedy equality with different batch size."""
137
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
138
- # NOTE Test is sensitive enough st if we don't enable chunked prefill
139
- # scheduling on baseline too, we get slightly different logprobs, ending
140
- # up sampling different tokens at the tail (ie top tokens don't change).
141
- # TL;DR: sd+cp == org+cp but sd+cp != org..is this expected?
142
- maybe_enable_chunked_prefill(prefill_chunk_size, baseline_llm_kwargs)
143
- run_equality_correctness_test(
144
- vllm_runner,
145
- common_llm_kwargs,
146
- per_test_common_llm_kwargs,
147
- baseline_llm_kwargs,
148
- test_llm_kwargs,
149
- batch_size,
150
- max_output_len=output_len,
151
- seed=seed,
152
- temperature=0.0,
153
- logprobs=logprobs,
154
- prompt_logprobs=logprobs,
155
- disable_logprobs=test_llm_kwargs["speculative_config"]
156
- ["disable_logprobs"])
157
-
158
-
159
-@pytest.mark.parametrize(
160
- "common_llm_kwargs",
161
- [{
162
- # Skip cuda graph recording for fast test.
163
- "enforce_eager": True,
164
-
165
- # Print spec metrics.
166
- "disable_log_stats": False,
167
-
168
- # Precision
169
- "dtype": PRECISION,
170
-
171
- # Main model
172
- "model_name": MAIN_MODEL,
173
- }])
174
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
175
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
176
-@pytest.mark.parametrize("test_llm_kwargs", [
177
- {
178
- "speculative_config": {
179
- "model": SPEC_MODEL,
180
- },
181
- },
182
-])
183
-@pytest.mark.parametrize("output_len", [2048])
184
-@pytest.mark.parametrize("batch_size", [1, 32])
185
-@pytest.mark.parametrize("seed", [1])
186
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 4])
187
-def test_mlp_e2e_acceptance_rate(vllm_runner, common_llm_kwargs,
188
- per_test_common_llm_kwargs,
189
- baseline_llm_kwargs, test_llm_kwargs,
190
- batch_size: int, output_len: int,
191
- prefill_chunk_size: int, seed: int):
192
- """Verify acceptance rate with different batch size and large output
193
- length."""
194
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
195
- run_equality_correctness_test(vllm_runner,
196
- common_llm_kwargs,
197
- per_test_common_llm_kwargs,
198
- baseline_llm_kwargs,
199
- test_llm_kwargs,
200
- batch_size,
201
- max_output_len=output_len,
202
- temperature=0.0,
203
- seed=seed,
204
- expected_acceptance_rate=0.48)
205
-
206
-
207
-@pytest.mark.parametrize(
208
- "common_llm_kwargs",
209
- [{
210
- # Skip cuda graph recording for fast test.
211
- "enforce_eager": True,
212
-
213
- # Print spec metrics.
214
- "disable_log_stats": False,
215
-
216
- # Precision
217
- "dtype": PRECISION,
218
-
219
- # Main model
220
- "model_name": MAIN_MODEL,
221
-
222
- # Speculative config
223
- "speculative_config": {
224
- "model": SPEC_MODEL,
225
- },
226
- }])
227
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
228
-@pytest.mark.parametrize("baseline_llm_kwargs", [{"seed": 1}])
229
-@pytest.mark.parametrize("test_llm_kwargs", [{"seed": 5}])
230
-@pytest.mark.parametrize("output_len", [64])
231
-@pytest.mark.parametrize("batch_size", [1, 32])
232
-@pytest.mark.parametrize("temperature", [1.0])
233
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 4])
234
-@pytest.mark.parametrize("seed", [1])
235
-def test_mlp_e2e_seeded_correctness(vllm_runner, common_llm_kwargs,
236
- per_test_common_llm_kwargs,
237
- baseline_llm_kwargs, test_llm_kwargs,
238
- batch_size: int, output_len: int,
239
- temperature: float,
240
- prefill_chunk_size: int, seed: int):
241
- """Verify seeded runs produce the same output."""
242
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
243
- maybe_enable_chunked_prefill(prefill_chunk_size, baseline_llm_kwargs)
244
- run_equality_correctness_test(vllm_runner,
245
- common_llm_kwargs,
246
- per_test_common_llm_kwargs,
247
- baseline_llm_kwargs,
248
- test_llm_kwargs,
249
- batch_size,
250
- max_output_len=output_len,
251
- temperature=temperature,
252
- seed=seed)
253
-
254
- # Ensure this same test does fail if we _don't_ include per-request seeds
255
- with pytest.raises(AssertionError):
256
- run_equality_correctness_test(vllm_runner,
257
- common_llm_kwargs,
258
- per_test_common_llm_kwargs,
259
- baseline_llm_kwargs,
260
- test_llm_kwargs,
261
- batch_size,
262
- max_output_len=output_len,
263
- temperature=temperature,
264
- seed=seed,
265
- disable_seed=True)
266
-
267
-
268
-@pytest.mark.parametrize(
269
- "common_llm_kwargs",
270
- [{
271
- "block_size": 16,
272
- # 2 for small prompt, 256//8 for generated.
273
- "num_gpu_blocks_override": 2 + 256 // 8,
274
- "max_model_len": (2 + 256 // 8) * 8,
275
-
276
- # Skip cuda graph recording for fast test.
277
- "enforce_eager": True,
278
-
279
- # Precision
280
- "dtype": PRECISION,
281
-
282
- # Main model
283
- "model_name": MAIN_MODEL,
284
- }])
285
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
286
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
287
-@pytest.mark.parametrize("test_llm_kwargs", [
288
- {
289
- "speculative_config": {
290
- "model": SPEC_MODEL,
291
- },
292
- },
293
-])
294
-@pytest.mark.parametrize(
295
- "output_len",
296
- [
297
- # Use small output len for fast test.
298
- 128,
299
- ])
300
-@pytest.mark.parametrize("batch_size", [4])
301
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 4])
302
-@pytest.mark.parametrize("seed", [1])
303
-def test_mlp_e2e_greedy_correctness_with_preemption(
304
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
305
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
306
- prefill_chunk_size: int, seed: int):
307
- """Verify greedy equality, even when some sequences are preempted mid-
308
- generation.
309
- """
310
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
311
- run_equality_correctness_test(vllm_runner,
312
- common_llm_kwargs,
313
- per_test_common_llm_kwargs,
314
- baseline_llm_kwargs,
315
- test_llm_kwargs,
316
- batch_size,
317
- max_output_len=output_len,
318
- seed=seed,
319
- temperature=0.0)
320
-
321
-
322
-@pytest.mark.parametrize(
323
- "common_llm_kwargs",
324
- [{
325
- "block_size": 16,
326
- # 2 for small prompt, 256//8 for generated.
327
- "num_gpu_blocks_override": 2 + 256 // 8,
328
- "max_model_len": (2 + 256 // 8) * 8,
329
-
330
- # Skip cuda graph recording for fast test.
331
- "enforce_eager": True,
332
-
333
- # Precision
334
- "dtype": PRECISION,
335
-
336
- # Main model
337
- "model_name": MAIN_MODEL,
338
- }])
339
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
340
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
341
-@pytest.mark.parametrize("test_llm_kwargs", [
342
- {
343
- "speculative_config": {
344
- "model": SPEC_MODEL,
345
- },
346
- },
347
-])
348
-@pytest.mark.parametrize(
349
- "output_len",
350
- [
351
- # Use small output len for fast test.
352
- 128,
353
- ])
354
-@pytest.mark.parametrize("batch_size", [4])
355
-@pytest.mark.parametrize("seed", [1])
356
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 4])
357
-def test_mlp_e2e_greedy_correctness_with_padding(
358
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
359
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
360
- prefill_chunk_size: int, seed: int):
361
- """Verify greedy equality when the vocab dimension is padded
362
- """
363
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
364
-
365
- # Default pad_to is 64, test model has vocab_size of 32000
366
- def patched_pad_vocab_size(vocab_size, pad_to=None):
367
- return pad_vocab_size(vocab_size, pad_to=32064)
368
-
369
- with patch(
370
- "vllm.model_executor.layers.vocab_parallel_embedding.pad_vocab_size",
371
- patched_pad_vocab_size):
372
- run_equality_correctness_test(vllm_runner,
373
- common_llm_kwargs,
374
- per_test_common_llm_kwargs,
375
- baseline_llm_kwargs,
376
- test_llm_kwargs,
377
- batch_size,
378
- max_output_len=output_len,
379
- seed=seed,
380
- temperature=0.0)
381
-
382
-
383
-@pytest.mark.parametrize(
384
- "common_llm_kwargs",
385
- [{
386
- # Skip cuda graph recording for fast test.
387
- "enforce_eager": True,
388
-
389
- # Precision
390
- "dtype": PRECISION,
391
-
392
- # Main model
393
- "model_name": MAIN_MODEL,
394
- }])
395
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
396
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
397
-@pytest.mark.parametrize(
398
- "test_llm_kwargs",
399
- [
400
- {
401
- "speculative_config": {
402
- "model": SPEC_MODEL,
403
- "num_speculative_tokens": k,
404
- },
405
- }
406
- # Try a range of num. speculative tokens
407
- for k in range(1, 1 + MAX_SPEC_TOKENS)
408
- ])
409
-@pytest.mark.parametrize("batch_size", [2])
410
-@pytest.mark.parametrize(
411
- "output_len",
412
- [
413
- # Use smaller output len for fast test.
414
- 32,
415
- ])
416
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 4])
417
-@pytest.mark.parametrize("seed", [1])
418
-def test_mlp_different_k(vllm_runner, common_llm_kwargs,
419
- per_test_common_llm_kwargs, baseline_llm_kwargs,
420
- test_llm_kwargs, batch_size: int,
421
- prefill_chunk_size: int, seed: int, output_len: int):
422
- """Verify that mlp speculative decoding produces exact equality
423
- to without spec decode with different values of num_speculative_tokens.
424
- """
425
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
426
- run_equality_correctness_test(vllm_runner,
427
- common_llm_kwargs,
428
- per_test_common_llm_kwargs,
429
- baseline_llm_kwargs,
430
- test_llm_kwargs,
431
- batch_size,
432
- max_output_len=output_len,
433
- seed=seed,
434
- temperature=0.0)
435
-
436
-
437
-@pytest.mark.parametrize(
438
- "common_llm_kwargs",
439
- [{
440
- # Skip cuda graph recording for fast test.
441
- "enforce_eager": True,
442
-
443
- # Precision
444
- "dtype": PRECISION,
445
-
446
- # Main model
447
- "model_name": MAIN_MODEL,
448
- }])
449
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
450
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
451
-@pytest.mark.parametrize("test_llm_kwargs", [{
452
- "speculative_config": {
453
- "model": SPEC_MODEL,
454
- "disable_by_batch_size": 4,
455
- },
456
-}])
457
-@pytest.mark.parametrize("batch_size", [1, 5])
458
-@pytest.mark.parametrize(
459
- "output_len",
460
- [
461
- # Use smaller output len for fast test.
462
- 32,
463
- ])
464
-# Speculative decoding is disabled when sequences reach decoding and the batch
465
-# consists of single-token requests. Hence we set `max_num_seqs`
466
-# >= `speculative_disable_by_batch_size` to test feature interaction.
467
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 4])
468
-@pytest.mark.parametrize("seed", [1])
469
-def test_mlp_disable_queue(vllm_runner, common_llm_kwargs,
470
- per_test_common_llm_kwargs, baseline_llm_kwargs,
471
- test_llm_kwargs, batch_size: int,
472
- prefill_chunk_size: int, seed: int,
473
- output_len: int):
474
- """Verify that mlp speculative decoding produces exact equality
475
- to without spec decode when speculation is disabled for large
476
- batch sizes.
477
- """
478
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
479
- run_equality_correctness_test(vllm_runner,
480
- common_llm_kwargs,
481
- per_test_common_llm_kwargs,
482
- baseline_llm_kwargs,
483
- test_llm_kwargs,
484
- batch_size,
485
- max_output_len=output_len,
486
- seed=seed,
487
- temperature=0.0)
488
-
489
-
490
-@pytest.mark.parametrize(
491
- "common_llm_kwargs",
492
- [{
493
- "model_name": MAIN_MODEL,
494
-
495
- # Skip cuda graph recording for fast test.
496
- "enforce_eager": True,
497
-
498
- # Precision
499
- "dtype": PRECISION,
500
- }])
501
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
502
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
503
-@pytest.mark.parametrize("test_llm_kwargs", [{
504
- "speculative_config": {
505
- "model": SPEC_MODEL,
506
- "disable_mqa_scorer": True,
507
- },
508
-}])
509
-@pytest.mark.parametrize("batch_size", [1, 5])
510
-@pytest.mark.parametrize(
511
- "output_len",
512
- [
513
- # Use smaller output len for fast test.
514
- 32,
515
- ])
516
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 4])
517
-@pytest.mark.parametrize("seed", [1])
518
-def test_mqa_scorer(vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
519
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int,
520
- output_len: int, prefill_chunk_size: int, seed: int):
521
- """Verify that speculative decoding generates the same output
522
- with batch expansion scorer and mqa scorer.
523
- """
524
- maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
525
- run_equality_correctness_test(vllm_runner,
526
- common_llm_kwargs,
527
- per_test_common_llm_kwargs,
528
- baseline_llm_kwargs,
529
- test_llm_kwargs,
530
- batch_size,
531
- max_output_len=output_len,
532
- seed=seed,
533
- temperature=0.0)
tests/spec_decode/e2e/test_mtp_correctness.py
deleted
-333
@@ -1,333 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""This docstring details important information on the testing methodology.
4
-
5
-Most of the tests rely on "greedy equality", where we expect the output of
6
-speculative decoding on a sequence to exactly match the output of normal non-
7
-speculative decoding.
8
-
9
-Since speculative decoding with rejection sampling guarantees that the output
10
-distribution matches the target model's output distribution (up to hardware
11
-numerics, see https://arxiv.org/pdf/2302.01318.pdf), we can expect greedy
12
-equality.
13
-
14
-However, we still need to verify below scenario could be passed:
15
- * Batch size 1 greedy equality
16
- * Batch size >1 greedy equality
17
- * Test greedy equality under preemption
18
- * Test greedy equality under various number of speculative tokens.
19
-
20
-With those tests, we can say at least, mtp would not break the
21
-correctness for the target model outputs.
22
-"""
23
-
24
-import pytest
25
-
26
-from .conftest import run_equality_correctness_test
27
-
28
-# main model
29
-MAIN_MODEL = "luccafong/deepseek_mtp_main_random"
30
-
31
-# max. number of speculative tokens: this corresponds to
32
-# num_nextn_predict_layers in the config.json of the speculator model.
33
-MAX_SPEC_TOKENS = 1
34
-
35
-# precision
36
-PRECISION = "bfloat16"
37
-
38
-
39
-@pytest.mark.parametrize(
40
- "common_llm_kwargs",
41
- [{
42
- # Skip cuda graph recording for fast test.
43
- "enforce_eager": True,
44
-
45
- # Print spec metrics.
46
- "disable_log_stats": False,
47
-
48
- # Precision
49
- "dtype": PRECISION,
50
-
51
- # Main model
52
- "model_name": MAIN_MODEL,
53
-
54
- # GPU memory utilization
55
- "gpu_memory_utilization": 0.85
56
- }])
57
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
58
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
59
-@pytest.mark.parametrize("test_llm_kwargs", [
60
- {
61
- "speculative_config": {
62
- "num_speculative_tokens": MAX_SPEC_TOKENS,
63
- },
64
- },
65
-])
66
-@pytest.mark.parametrize("output_len", [
67
- 128,
68
-])
69
-@pytest.mark.parametrize("batch_size", [1, 32])
70
-@pytest.mark.parametrize("seed", [1])
71
-def test_mtp_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
72
- per_test_common_llm_kwargs,
73
- baseline_llm_kwargs, test_llm_kwargs,
74
- batch_size: int, output_len: int,
75
- seed: int):
76
-
77
- run_equality_correctness_test(vllm_runner, common_llm_kwargs,
78
- per_test_common_llm_kwargs,
79
- baseline_llm_kwargs, test_llm_kwargs,
80
- batch_size, output_len, seed)
81
-
82
-
83
-@pytest.mark.parametrize(
84
- "common_llm_kwargs",
85
- [{
86
- # Skip cuda graph recording for fast test.
87
- "enforce_eager": True,
88
-
89
- # Print spec metrics.
90
- "disable_log_stats": False,
91
-
92
- # Precision
93
- "dtype": PRECISION,
94
-
95
- # Main model
96
- "model_name": MAIN_MODEL,
97
-
98
- # GPU memory utilization
99
- "gpu_memory_utilization": 0.85
100
- }])
101
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
102
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
103
-@pytest.mark.parametrize("test_llm_kwargs", [
104
- {
105
- "speculative_config": {
106
- "num_speculative_tokens": MAX_SPEC_TOKENS,
107
- "disable_logprobs": False,
108
- },
109
- },
110
- {
111
- "speculative_config": {
112
- "num_speculative_tokens": MAX_SPEC_TOKENS,
113
- "disable_logprobs": True,
114
- },
115
- },
116
-])
117
-@pytest.mark.parametrize("output_len", [
118
- 128,
119
-])
120
-@pytest.mark.parametrize("batch_size", [8])
121
-@pytest.mark.parametrize("seed", [1])
122
-@pytest.mark.parametrize("logprobs", [1, 6])
123
-def test_mtp_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
124
- per_test_common_llm_kwargs,
125
- baseline_llm_kwargs, test_llm_kwargs,
126
- batch_size: int, output_len: int, seed: int,
127
- logprobs: int):
128
-
129
- run_equality_correctness_test(
130
- vllm_runner,
131
- common_llm_kwargs,
132
- per_test_common_llm_kwargs,
133
- baseline_llm_kwargs,
134
- test_llm_kwargs,
135
- batch_size,
136
- output_len,
137
- seed,
138
- logprobs=logprobs,
139
- prompt_logprobs=logprobs,
140
- disable_logprobs=test_llm_kwargs["speculative_config"]
141
- ["disable_logprobs"])
142
-
143
-
144
-@pytest.mark.parametrize(
145
- "common_llm_kwargs",
146
- [{
147
- "enforce_eager": False,
148
-
149
- # Print spec metrics.
150
- "disable_log_stats": False,
151
-
152
- # Precision
153
- "dtype": PRECISION,
154
-
155
- # Main model
156
- "model_name": MAIN_MODEL,
157
- "gpu_memory_utilization": 0.85
158
- }])
159
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
160
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
161
-@pytest.mark.parametrize("test_llm_kwargs", [
162
- {
163
- "speculative_config": {
164
- "num_speculative_tokens": MAX_SPEC_TOKENS,
165
- },
166
- },
167
-])
168
-@pytest.mark.parametrize("output_len", [
169
- 128,
170
-])
171
-@pytest.mark.parametrize("batch_size", [1, 32])
172
-@pytest.mark.parametrize("seed", [1])
173
-def test_mtp_e2e_greedy_correctness_cuda_graph(vllm_runner, common_llm_kwargs,
174
- per_test_common_llm_kwargs,
175
- baseline_llm_kwargs,
176
- test_llm_kwargs,
177
- batch_size: int,
178
- output_len: int, seed: int):
179
- """Verify greedy equality with cuda graph enabled and different
180
- batch sizes."""
181
- run_equality_correctness_test(vllm_runner, common_llm_kwargs,
182
- per_test_common_llm_kwargs,
183
- baseline_llm_kwargs, test_llm_kwargs,
184
- batch_size, output_len, seed)
185
-
186
-
187
-@pytest.mark.parametrize(
188
- "common_llm_kwargs",
189
- [{
190
- "block_size": 8,
191
- # 2 for small prompt, 256//8 for generated.
192
- "num_gpu_blocks_override": 2 + 256 // 8,
193
- "max_model_len": (2 + 256 // 8) * 8,
194
-
195
- # Skip cuda graph recording for fast test.
196
- "enforce_eager": True,
197
-
198
- # Precision
199
- "dtype": PRECISION,
200
-
201
- # Main model
202
- "model_name": MAIN_MODEL,
203
-
204
- # GPU memory utilization
205
- "gpu_memory_utilization": 0.9
206
- }])
207
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
208
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
209
-@pytest.mark.parametrize("test_llm_kwargs", [
210
- {
211
- "speculative_config": {
212
- "num_speculative_tokens": MAX_SPEC_TOKENS,
213
- },
214
- },
215
-])
216
-@pytest.mark.parametrize(
217
- "output_len",
218
- [
219
- # Use small output len for fast test.
220
- 128,
221
- ])
222
-@pytest.mark.parametrize("batch_size", [4])
223
-@pytest.mark.parametrize("seed", [1])
224
-def test_mtp_e2e_greedy_correctness_with_preemption(
225
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
226
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
227
- seed: int):
228
- """Verify greedy equality, even when some sequences are preempted mid-
229
- generation.
230
- """
231
- run_equality_correctness_test(vllm_runner, common_llm_kwargs,
232
- per_test_common_llm_kwargs,
233
- baseline_llm_kwargs, test_llm_kwargs,
234
- batch_size, output_len, seed)
235
-
236
-
237
-@pytest.mark.parametrize(
238
- "common_llm_kwargs",
239
- [{
240
- # Skip cuda graph recording for fast test.
241
- "enforce_eager": True,
242
-
243
- # Precision
244
- "dtype": PRECISION,
245
-
246
- # Main model
247
- "model_name": MAIN_MODEL,
248
-
249
- # GPU memory utilization
250
- "gpu_memory_utilization": 0.9
251
- }])
252
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
253
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
254
-@pytest.mark.parametrize(
255
- "test_llm_kwargs",
256
- [
257
- {
258
- "speculative_config": {
259
- "num_speculative_tokens": k,
260
- },
261
- }
262
- # Try a range of num. speculative tokens
263
- for k in range(1, 1 + MAX_SPEC_TOKENS)
264
- ])
265
-@pytest.mark.parametrize("batch_size", [2])
266
-@pytest.mark.parametrize(
267
- "output_len",
268
- [
269
- # Use smaller output len for fast test.
270
- 32,
271
- ])
272
-@pytest.mark.parametrize("seed", [1])
273
-def test_mtp_different_k(vllm_runner, common_llm_kwargs,
274
- per_test_common_llm_kwargs, baseline_llm_kwargs,
275
- test_llm_kwargs, batch_size: int, output_len: int,
276
- seed: int):
277
- """Verify that mtp speculative decoding produces exact equality
278
- to without spec decode with different values of num_speculative_tokens.
279
- """
280
- run_equality_correctness_test(vllm_runner, common_llm_kwargs,
281
- per_test_common_llm_kwargs,
282
- baseline_llm_kwargs, test_llm_kwargs,
283
- batch_size, output_len, seed)
284
-
285
-
286
-@pytest.mark.parametrize(
287
- "common_llm_kwargs",
288
- [{
289
- # Skip cuda graph recording for fast test.
290
- "enforce_eager": True,
291
-
292
- # Precision
293
- "dtype": PRECISION,
294
-
295
- # Main model
296
- "model_name": MAIN_MODEL,
297
-
298
- # GPU memory utilization
299
- "gpu_memory_utilization": 0.9
300
- }])
301
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
302
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
303
-@pytest.mark.parametrize("test_llm_kwargs", [{
304
- "speculative_config": {
305
- "num_speculative_tokens": MAX_SPEC_TOKENS,
306
- "disable_by_batch_size": 4
307
- },
308
-}])
309
-@pytest.mark.parametrize("batch_size", [1, 5])
310
-@pytest.mark.parametrize(
311
- "output_len",
312
- [
313
- # Use smaller output len for fast test.
314
- 32,
315
- ])
316
-@pytest.mark.parametrize("seed", [1])
317
-def test_mtp_disable_queue(vllm_runner, common_llm_kwargs,
318
- per_test_common_llm_kwargs, baseline_llm_kwargs,
319
- test_llm_kwargs, batch_size: int, output_len: int,
320
- seed: int):
321
- """Verify that mtp speculative decoding produces exact equality
322
- to without spec decode when speculation is disabled for large
323
- batch sizes.
324
- """
325
- run_equality_correctness_test(vllm_runner, common_llm_kwargs,
326
- per_test_common_llm_kwargs,
327
- baseline_llm_kwargs, test_llm_kwargs,
328
- batch_size, output_len, seed)
329
-
330
-
331
-if __name__ == "__main__":
332
- import pytest
333
- pytest.main([__file__])
tests/spec_decode/e2e/test_multistep_correctness.py
deleted
-842
@@ -1,842 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""The tests in this file verify end-to-end speculative decoding correctness.
4
-
5
-This docstring details important information on the testing methodology.
6
-
7
-Most of the tests rely on "greedy equality", where we expect the output of
8
-speculative decoding on a sequence to exactly match the output of normal non-
9
-speculative decoding.
10
-
11
-Since speculative decoding with rejection sampling guarantees that the output
12
-distribution matches the target model's output distribution (up to hardware
13
-numerics, see https://arxiv.org/pdf/2302.01318.pdf), we can expect greedy
14
-equality. This gives us good coverage of temp=0.
15
-
16
-At temp=0, the TypicalAcceptanceSampler ensures that only the tokens with the
17
-highest probability in the target distribution are accepted. Therefore, we can
18
-expect greedy equality for the TypicalAcceptanceSampler at temp=0.
19
-
20
-For temp>0, we rely on unit tests on the rejection sampler to verify that the
21
-output distribution is the same with spec decode vs. no spec decode (this would
22
-be prohibitively expensive to run with a real model). Similarly, for the
23
-TypicalAcceptance sampler also, we rely on unit tests to validate temp>0
24
-test cases.
25
-
26
-NOTE: Speculative decoding's distribution equality requires that the measured
27
-distributions of the target model and proposal model be deterministic given the
28
-same input. vLLM largely guarantees this.
29
-
30
-@cadedaniel has seen cases where the output probabilities of a draft/target
31
-model change slightly with certain batch sizes or prompts, even with Torch
32
-determinism flags set. It is unclear if this is a bug in vLLM, due to non-
33
-determinism in on-device batched operations, a bug in vLLM's spec decode
34
-implementation, or the "hardware numerics" limitations. Either way, rejection
35
-sampling ensures the output distribution matches the target model, but it breaks
36
-greedy-equality tests for those batch sizes/prompts.
37
-"""
38
-
39
-from itertools import cycle
40
-
41
-import pytest
42
-from transformers import AutoTokenizer
43
-
44
-from vllm import SamplingParams
45
-
46
-from ...utils import create_new_process_for_each_test
47
-from .conftest import (get_output_from_llm_generator,
48
- run_equality_correctness_test)
49
-
50
-
51
-@pytest.mark.parametrize(
52
- "common_llm_kwargs",
53
- [{
54
- # Use a small model for a fast test.
55
- # Note this is repeated in the test body; to initialize a tokenizer.
56
- "model": "JackFram/llama-68m",
57
-
58
- # Skip cuda graph recording for fast test.
59
- "enforce_eager": True,
60
-
61
- # The original model is float32, keep it for numerical stability.
62
- "dtype": "float32",
63
- }])
64
-@pytest.mark.parametrize(
65
- "per_test_common_llm_kwargs",
66
- [
67
- {
68
- "speculative_config": {
69
- "model": "JackFram/llama-68m",
70
- "num_speculative_tokens": 5,
71
- },
72
- "enable_chunked_prefill": False,
73
- },
74
- {
75
- # Chunked prefill enabled with small value
76
- # to make sure we get mixed batches.
77
- "speculative_config": {
78
- "model": "JackFram/llama-68m",
79
- "num_speculative_tokens": 5,
80
- },
81
- "enable_chunked_prefill": True,
82
- "max_num_batched_tokens": 4,
83
- "max_num_seqs": 4
84
- },
85
- {
86
- # Verify the detokenizer assertions in the test work when spec
87
- # decode is disabled.
88
- },
89
- ])
90
-@pytest.mark.parametrize("test_llm_kwargs", [{}])
91
-@pytest.mark.parametrize("batch_size", [1, 32])
92
-@pytest.mark.parametrize("seed", [1])
93
-@create_new_process_for_each_test()
94
-def test_spec_decode_e2e_with_detokenization(test_llm_generator,
95
- batch_size: int):
96
- """Run generation with speculative decoding on a batch. Verify the engine
97
- generates the correct number of tokens (via ignore_eos=True), and that the
98
- detokenization matches HF transformers.
99
- """
100
- output_len = 32
101
- temperature = 0.0
102
-
103
- prompts = [
104
- "Hello, my name is",
105
- "The president of the United States is",
106
- "The capital of France is",
107
- "The future of AI is",
108
- ]
109
-
110
- prompts = [prompt for prompt, _ in zip(cycle(prompts), range(batch_size))]
111
-
112
- sampling_params = SamplingParams(
113
- max_tokens=output_len,
114
- ignore_eos=True,
115
- temperature=temperature,
116
- )
117
-
118
- batch_tokens, batch_token_ids, _ = get_output_from_llm_generator(
119
- test_llm_generator, prompts, sampling_params)
120
-
121
- # Expect a generation for each prompt in the batch.
122
- assert len(batch_token_ids) == len(prompts)
123
-
124
- # Expect each generation to have expected number of tokens (note ignore_eos
125
- # is True).
126
- assert [len(token_ids)
127
- for token_ids in batch_token_ids] == ([output_len] * batch_size)
128
-
129
- # Expect detokenized string to match.
130
- tok = AutoTokenizer.from_pretrained("JackFram/llama-68m")
131
- for actual_tokens, actual_token_ids in zip(batch_tokens, batch_token_ids):
132
- expected_tokens = tok.decode(actual_token_ids)
133
- print(f"{actual_token_ids=}")
134
- assert actual_tokens.strip() == expected_tokens.strip()
135
-
136
-
137
-@pytest.mark.parametrize(
138
- "common_llm_kwargs",
139
- [{
140
- # Skip cuda graph recording for fast test.
141
- "enforce_eager": True,
142
-
143
- # Print spec metrics.
144
- "disable_log_stats": False,
145
-
146
- # The original model is float32, keep it for numerical stability.
147
- "dtype": "float32",
148
- }])
149
-@pytest.mark.parametrize(
150
- "per_test_common_llm_kwargs",
151
- [
152
- # Try two different tiny base models.
153
- # Note that one is equal to the draft model, another isn't.
154
- {
155
- "model_name": "JackFram/llama-68m",
156
- },
157
- {
158
- "model_name": "JackFram/llama-160m",
159
- },
160
- ])
161
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
162
-@pytest.mark.parametrize("test_llm_kwargs", [{
163
- "speculative_config": {
164
- "model": "JackFram/llama-68m",
165
- "num_speculative_tokens": 5,
166
- "disable_logprobs": False,
167
- },
168
- "enable_chunked_prefill": False,
169
-}, {
170
- "speculative_config": {
171
- "model": "JackFram/llama-68m",
172
- "num_speculative_tokens": 3,
173
- "disable_logprobs": False,
174
- },
175
- "enable_chunked_prefill": True,
176
- "max_num_batched_tokens": 4,
177
- "max_num_seqs": 4,
178
-}])
179
-@pytest.mark.parametrize(
180
- "output_len",
181
- [
182
- # Use long output len for the small model test.
183
- 10,
184
- ])
185
-@pytest.mark.parametrize("batch_size", [1])
186
-@pytest.mark.parametrize("seed", [1])
187
-@create_new_process_for_each_test()
188
-def test_spec_decode_e2e_greedy_correctness_tiny_model_bs1(
189
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
190
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
191
- seed: int):
192
- """Verify greedy equality on a tiny model with batch size of one.
193
-
194
- Since this test is cheaper than other e2e correctness tests, we generate
195
- with a higher output_len.
196
-
197
- When the draft model is the same as the target model, we further check
198
- whether all speculative tokens are accepted.
199
- """
200
- ensure_all_accepted = per_test_common_llm_kwargs.get(
201
- "model_name") == test_llm_kwargs.get("speculative_config")["model"]
202
- run_equality_correctness_test(vllm_runner,
203
- common_llm_kwargs,
204
- per_test_common_llm_kwargs,
205
- baseline_llm_kwargs,
206
- test_llm_kwargs,
207
- batch_size,
208
- max_output_len=output_len,
209
- seed=seed,
210
- prompt_logprobs=2,
211
- logprobs=2,
212
- disable_logprobs=False,
213
- temperature=0.0,
214
- ensure_all_accepted=ensure_all_accepted)
215
-
216
-
217
-@pytest.mark.parametrize(
218
- "common_llm_kwargs",
219
- [{
220
- # Skip cuda graph recording for fast test.
221
- "enforce_eager": True,
222
-
223
- # Print spec metrics.
224
- "disable_log_stats": False,
225
-
226
- # The original model is float32, keep it for numerical stability.
227
- "dtype": "float32",
228
- }])
229
-@pytest.mark.parametrize(
230
- "per_test_common_llm_kwargs",
231
- [
232
- # Try two different tiny base models.
233
- # Note that one is equal to the draft model, another isn't.
234
- {
235
- "model_name": "JackFram/llama-68m",
236
- },
237
- {
238
- "model_name": "JackFram/llama-160m",
239
- },
240
- ])
241
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
242
-@pytest.mark.parametrize("test_llm_kwargs", [
243
- {
244
- "speculative_config": {
245
- "model": "JackFram/llama-68m",
246
- "num_speculative_tokens": 5,
247
- },
248
- "enable_chunked_prefill": False,
249
- },
250
- {
251
- "speculative_config": {
252
- "model": "JackFram/llama-68m",
253
- "num_speculative_tokens": 5,
254
- },
255
- "enable_chunked_prefill": True,
256
- "max_num_batched_tokens": 4,
257
- "max_num_seqs": 4
258
- },
259
-])
260
-@pytest.mark.parametrize(
261
- "output_len",
262
- [
263
- # Use small output len for fast test.
264
- 256,
265
- ])
266
-@pytest.mark.parametrize("batch_size", [64])
267
-@pytest.mark.parametrize("seed", [1])
268
-@create_new_process_for_each_test()
269
-def test_spec_decode_e2e_greedy_correctness_tiny_model_large_bs(
270
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
271
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
272
- seed: int):
273
- """Verify greedy equality on a tiny model and large batch size.
274
- """
275
- run_equality_correctness_test(vllm_runner,
276
- common_llm_kwargs,
277
- per_test_common_llm_kwargs,
278
- baseline_llm_kwargs,
279
- test_llm_kwargs,
280
- batch_size,
281
- max_output_len=output_len,
282
- seed=seed,
283
- temperature=0.0)
284
-
285
-
286
-@pytest.mark.parametrize(
287
- "common_llm_kwargs",
288
- [{
289
- # Skip cuda graph recording for fast test.
290
- "enforce_eager": True,
291
-
292
- # The original model is float32, keep it for numerical stability.
293
- "dtype": "float32",
294
- }])
295
-@pytest.mark.parametrize(
296
- "per_test_common_llm_kwargs",
297
- [
298
- # Try two different tiny base models.
299
- # Note that one is equal to the draft model, another isn't.
300
- {
301
- "model_name": "JackFram/llama-68m",
302
- },
303
- {
304
- "model_name": "JackFram/llama-160m",
305
- },
306
- ])
307
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
308
-@pytest.mark.parametrize("test_llm_kwargs", [
309
- {
310
- "speculative_config": {
311
- "model": "JackFram/llama-68m",
312
- "num_speculative_tokens": 5,
313
- },
314
- "enable_chunked_prefill": False,
315
- },
316
- {
317
- "speculative_config": {
318
- "model": "JackFram/llama-68m",
319
- "num_speculative_tokens": 5,
320
- },
321
- "enable_chunked_prefill": True,
322
- "max_num_batched_tokens": 4,
323
- "max_num_seqs": 4
324
- },
325
-])
326
-@pytest.mark.parametrize("max_output_len", [
327
- 256,
328
-])
329
-@pytest.mark.parametrize("batch_size", [32])
330
-@pytest.mark.parametrize("seed", [1])
331
-@create_new_process_for_each_test()
332
-def test_spec_decode_e2e_greedy_correctness_tiny_model_large_bs_diff_output_len(
333
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
334
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int,
335
- max_output_len: int, seed: int):
336
- """Verify greedy equality on a tiny model, with a large batch size, and when
337
- sampling respects the EOS token.
338
- """
339
- run_equality_correctness_test(vllm_runner,
340
- common_llm_kwargs,
341
- per_test_common_llm_kwargs,
342
- baseline_llm_kwargs,
343
- test_llm_kwargs,
344
- batch_size,
345
- max_output_len,
346
- seed=seed,
347
- temperature=0.0,
348
- ignore_eos=False)
349
-
350
-
351
-@pytest.mark.parametrize(
352
- "common_llm_kwargs",
353
- [{
354
- # A "real" model (not tiny).
355
- "model_name": "meta-llama/Llama-2-7b-chat-hf",
356
-
357
- # Skip cuda graph recording for fast test.
358
- "enforce_eager": True,
359
-
360
- # Print spec metrics.
361
- "disable_log_stats": False,
362
- }])
363
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
364
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
365
-@pytest.mark.parametrize("test_llm_kwargs", [
366
- {
367
- "speculative_config": {
368
- "model": "JackFram/llama-68m",
369
- "num_speculative_tokens": 5,
370
- },
371
- "enable_chunked_prefill": False,
372
- },
373
- {
374
- "speculative_config": {
375
- "model": "JackFram/llama-68m",
376
- "num_speculative_tokens": 5,
377
- },
378
- "enable_chunked_prefill": True,
379
- "max_num_batched_tokens": 4,
380
- "max_num_seqs": 4
381
- },
382
-])
383
-@pytest.mark.parametrize("batch_size", [1])
384
-@pytest.mark.parametrize(
385
- "output_len",
386
- [
387
- # Use decently long output len for a high quality test.
388
- 256,
389
- ])
390
-@pytest.mark.parametrize("seed", [1])
391
-@create_new_process_for_each_test()
392
-def test_spec_decode_e2e_greedy_correctness_real_model_bs1(
393
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
394
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
395
- seed: int):
396
- """Verify greedy equality on a "real" model and batch size of 1. This is
397
- separate from large BS tests to make identifying the source of bugs easier.
398
- """
399
- run_equality_correctness_test(vllm_runner,
400
- common_llm_kwargs,
401
- per_test_common_llm_kwargs,
402
- baseline_llm_kwargs,
403
- test_llm_kwargs,
404
- batch_size,
405
- max_output_len=output_len,
406
- seed=seed,
407
- temperature=0.0)
408
-
409
-
410
-@pytest.mark.parametrize(
411
- "common_llm_kwargs",
412
- [{
413
- # A "real" model (not tiny).
414
- "model_name": "meta-llama/Llama-2-7b-chat-hf",
415
-
416
- # Skip cuda graph recording for fast test.
417
- "enforce_eager": True,
418
-
419
- # Print spec metrics.
420
- "disable_log_stats": False,
421
- }])
422
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
423
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
424
-@pytest.mark.parametrize("test_llm_kwargs", [
425
- {
426
- "speculative_config": {
427
- "model": "JackFram/llama-68m",
428
- "num_speculative_tokens": 5,
429
- },
430
- "enable_chunked_prefill": False,
431
- },
432
- {
433
- "speculative_config": {
434
- "model": "JackFram/llama-68m",
435
- "num_speculative_tokens": 5,
436
- },
437
- "enable_chunked_prefill": True,
438
- "max_num_batched_tokens": 4,
439
- "max_num_seqs": 4
440
- },
441
-])
442
-@pytest.mark.parametrize("batch_size", [32])
443
-@pytest.mark.parametrize(
444
- "output_len",
445
- [
446
- # Use smaller output len for fast test.
447
- 64,
448
- ])
449
-@pytest.mark.parametrize("seed", [1])
450
-@create_new_process_for_each_test()
451
-def test_spec_decode_e2e_greedy_correctness_real_model_large_bs(
452
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
453
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
454
- seed: int):
455
- """Verify greedy equality with a "real" model on a nontrivial batch size.
456
- This is the closest test to a real production workload.
457
- """
458
- run_equality_correctness_test(vllm_runner,
459
- common_llm_kwargs,
460
- per_test_common_llm_kwargs,
461
- baseline_llm_kwargs,
462
- test_llm_kwargs,
463
- batch_size,
464
- max_output_len=output_len,
465
- seed=seed,
466
- temperature=0.0)
467
-
468
-
469
-@pytest.mark.parametrize(
470
- "common_llm_kwargs",
471
- [{
472
- "block_size": 16,
473
- # 2 for small prompt, 256//8 for generated.
474
- "num_gpu_blocks_override": 2 + 256 // 8,
475
- "max_model_len": (2 + 256 // 8) * 8,
476
-
477
- # Skip cuda graph recording for fast test.
478
- "enforce_eager": True,
479
- # The original model is float32, keep it for numerical stability.
480
- "dtype": "float32",
481
- }])
482
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [
483
- {
484
- "model_name": "JackFram/llama-160m",
485
- },
486
-])
487
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
488
-@pytest.mark.parametrize("test_llm_kwargs", [
489
- {
490
- "speculative_config": {
491
- "model": "JackFram/llama-68m",
492
- "num_speculative_tokens": 5,
493
- },
494
- "enable_chunked_prefill": False,
495
- },
496
- {
497
- "speculative_config": {
498
- "model": "JackFram/llama-68m",
499
- "num_speculative_tokens": 5,
500
- },
501
- "enable_chunked_prefill": True,
502
- "max_num_batched_tokens": 4,
503
- "max_num_seqs": 4
504
- },
505
-])
506
-@pytest.mark.parametrize(
507
- "output_len",
508
- [
509
- # Use small output len for fast test.
510
- 256,
511
- ])
512
-@pytest.mark.parametrize("batch_size", [4])
513
-@pytest.mark.parametrize("seed", [1])
514
-@create_new_process_for_each_test()
515
-def test_spec_decode_e2e_greedy_correctness_with_preemption(
516
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
517
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
518
- seed: int):
519
- """Verify greedy equality, even when some sequences are preempted mid-
520
- generation.
521
- """
522
- run_equality_correctness_test(vllm_runner,
523
- common_llm_kwargs,
524
- per_test_common_llm_kwargs,
525
- baseline_llm_kwargs,
526
- test_llm_kwargs,
527
- batch_size,
528
- max_output_len=output_len,
529
- seed=seed,
530
- temperature=0.0)
531
-
532
-
533
-@pytest.mark.parametrize(
534
- "common_llm_kwargs",
535
- [{
536
- "model_name": "JackFram/llama-160m",
537
-
538
- # Skip cuda graph recording for fast test.
539
- "enforce_eager": True,
540
- # The original model is float32, keep it for numerical stability.
541
- "dtype": "float32",
542
- }])
543
-@pytest.mark.parametrize(
544
- "per_test_common_llm_kwargs",
545
- [
546
- # https://github.com/triton-lang/triton/issues/2266 tl.dot
547
- # doesn't support embedding < 16
548
- {
549
- "block_size": 16,
550
- },
551
- {
552
- "block_size": 32,
553
- },
554
- ])
555
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
556
-@pytest.mark.parametrize("test_llm_kwargs", [
557
- {
558
- "speculative_config": {
559
- "model": "JackFram/llama-68m",
560
- "num_speculative_tokens": 5,
561
- },
562
- "enable_chunked_prefill": False,
563
- },
564
- {
565
- "speculative_config": {
566
- "model": "JackFram/llama-68m",
567
- "num_speculative_tokens": 5,
568
- },
569
- "enable_chunked_prefill": True,
570
- "max_num_batched_tokens": 4,
571
- "max_num_seqs": 4
572
- },
573
-])
574
-@pytest.mark.parametrize("batch_size", [2])
575
-@pytest.mark.parametrize(
576
- "output_len",
577
- [
578
- # Use smaller output len for fast test.
579
- 32,
580
- ])
581
-@pytest.mark.parametrize("seed", [1])
582
-@create_new_process_for_each_test()
583
-def test_spec_decode_different_block_size(vllm_runner, common_llm_kwargs,
584
- per_test_common_llm_kwargs,
585
- baseline_llm_kwargs, test_llm_kwargs,
586
- batch_size: int, output_len: int,
587
- seed: int):
588
- """Verify greedy equality over different block sizes.
589
- """
590
- run_equality_correctness_test(vllm_runner,
591
- common_llm_kwargs,
592
- per_test_common_llm_kwargs,
593
- baseline_llm_kwargs,
594
- test_llm_kwargs,
595
- batch_size,
596
- max_output_len=output_len,
597
- seed=seed,
598
- temperature=0.0)
599
-
600
-
601
-@pytest.mark.parametrize(
602
- "common_llm_kwargs",
603
- [{
604
- "model_name": "JackFram/llama-160m",
605
-
606
- # Skip cuda graph recording for fast test.
607
- "enforce_eager": True,
608
- # The original model is float32, keep it for numerical stability.
609
- "dtype": "float32",
610
- }])
611
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
612
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
613
-@pytest.mark.parametrize(
614
- "test_llm_kwargs",
615
- [
616
- {
617
-
618
- # Artificially limit the draft model max model len; this forces vLLM
619
- # to skip speculation once the sequences grow beyond 32-k tokens.
620
- "speculative_config": {
621
- "model": "JackFram/llama-68m",
622
- "num_speculative_tokens": 5,
623
- "max_model_len": 32,
624
- },
625
- "enable_chunked_prefill": False,
626
- },
627
- {
628
- "speculative_config": {
629
- "model": "JackFram/llama-68m",
630
- "num_speculative_tokens": 5,
631
- "max_model_len": 32,
632
- },
633
- "enable_chunked_prefill": True,
634
- "max_num_batched_tokens": 4,
635
- "max_num_seqs": 4,
636
- },
637
- ])
638
-@pytest.mark.parametrize("batch_size", [8])
639
-@pytest.mark.parametrize(
640
- "output_len",
641
- [
642
- # This must be a good bit larger than speculative_max_model_len so that
643
- # we can test the case where all seqs are skipped, but still small to
644
- # ensure fast test.
645
- 64,
646
- ])
647
-@pytest.mark.parametrize("seed", [1])
648
-@create_new_process_for_each_test()
649
-def test_skip_speculation(vllm_runner, common_llm_kwargs,
650
- per_test_common_llm_kwargs, baseline_llm_kwargs,
651
- test_llm_kwargs, batch_size: int, output_len: int,
652
- seed: int):
653
- """Verify greedy equality when some (or all) sequences skip speculation.
654
- We do this by setting the max model len of the draft model to an
655
- artificially low value, such that when the sequences grow beyond it, they
656
- are skipped in speculative decoding.
657
- """
658
- run_equality_correctness_test(vllm_runner,
659
- common_llm_kwargs,
660
- per_test_common_llm_kwargs,
661
- baseline_llm_kwargs,
662
- test_llm_kwargs,
663
- batch_size,
664
- max_output_len=output_len,
665
- seed=seed,
666
- temperature=0.0)
667
-
668
-
669
-@pytest.mark.parametrize(
670
- "common_llm_kwargs",
671
- [{
672
- "model_name": "JackFram/llama-160m",
673
-
674
- # Skip cuda graph recording for fast test.
675
- "enforce_eager": True,
676
- # The original model is float32, keep it for numerical stability.
677
- "dtype": "float32",
678
- }])
679
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
680
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
681
-@pytest.mark.parametrize("test_llm_kwargs", [
682
- {
683
- "speculative_config": {
684
- "model": "JackFram/llama-68m",
685
- "num_speculative_tokens": 5,
686
- "disable_by_batch_size": 2,
687
- },
688
- "enable_chunked_prefill": False,
689
- },
690
- {
691
- "speculative_config": {
692
- "model": "JackFram/llama-68m",
693
- "num_speculative_tokens": 5,
694
- "disable_by_batch_size": 2,
695
- },
696
- "enable_chunked_prefill": True,
697
- "max_num_batched_tokens": 4,
698
- "max_num_seqs": 4,
699
- },
700
-])
701
-@pytest.mark.parametrize("batch_size", [8])
702
-@pytest.mark.parametrize("output_len", [10])
703
-@pytest.mark.parametrize("seed", [1])
704
-@create_new_process_for_each_test()
705
-def test_disable_speculation(vllm_runner, common_llm_kwargs,
706
- per_test_common_llm_kwargs, baseline_llm_kwargs,
707
- test_llm_kwargs, batch_size: int, output_len: int,
708
- seed: int):
709
- """Verify greedy equality when all sequences disable speculation.
710
- """
711
- run_equality_correctness_test(vllm_runner,
712
- common_llm_kwargs,
713
- per_test_common_llm_kwargs,
714
- baseline_llm_kwargs,
715
- test_llm_kwargs,
716
- batch_size,
717
- max_output_len=output_len,
718
- seed=seed,
719
- temperature=0.0)
720
-
721
-
722
-@pytest.mark.parametrize(
723
- "common_llm_kwargs",
724
- [{
725
- "model_name": "JackFram/llama-68m",
726
-
727
- # Skip cuda graph recording for fast test.
728
- "enforce_eager": True,
729
- # The original model is float32, keep it for numerical stability.
730
- "dtype": "float32",
731
- }])
732
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
733
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
734
-@pytest.mark.parametrize(
735
- "test_llm_kwargs",
736
- [
737
- {
738
- "speculative_config": {
739
- "model": "JackFram/llama-68m",
740
- "num_speculative_tokens": k,
741
- },
742
- "enable_chunked_prefill": False,
743
- }
744
- # Try a range of common k, as well as large speculation.
745
- for k in [1, 2, 3, 4, 5, 6, 7, 8, 9, 63]
746
- ] + [{
747
- "speculative_config": {
748
- "model": "JackFram/llama-68m",
749
- "num_speculative_tokens": k,
750
- },
751
- "enable_chunked_prefill": True,
752
- "max_num_batched_tokens": 4,
753
- "max_num_seqs": 4,
754
- } for k in [1, 2, 3, 4, 5, 6, 7, 8, 9, 63]])
755
-@pytest.mark.parametrize("batch_size", [2])
756
-@pytest.mark.parametrize(
757
- "output_len",
758
- [
759
- # Use smaller output len for fast test.
760
- 32,
761
- ])
762
-@pytest.mark.parametrize("seed", [1])
763
-@create_new_process_for_each_test()
764
-def test_many_k(vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
765
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int,
766
- output_len: int, seed: int):
767
- """Verify that speculative decoding produces exact equality to without spec
768
- decode with many different values of k.
769
- """
770
- run_equality_correctness_test(vllm_runner,
771
- common_llm_kwargs,
772
- per_test_common_llm_kwargs,
773
- baseline_llm_kwargs,
774
- test_llm_kwargs,
775
- batch_size,
776
- max_output_len=output_len,
777
- seed=seed,
778
- temperature=0.0)
779
-
780
-
781
-@pytest.mark.parametrize(
782
- "common_llm_kwargs",
783
- [{
784
- "model_name": "JackFram/llama-160m",
785
-
786
- # Skip cuda graph recording for fast test.
787
- "enforce_eager": True,
788
- # The original model is float32, keep it for numerical stability.
789
- "dtype": "float32",
790
- }])
791
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
792
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
793
-@pytest.mark.parametrize(
794
- "test_llm_kwargs",
795
- [
796
- {
797
- "speculative_config": {
798
- "model": "JackFram/llama-68m",
799
- "num_speculative_tokens": k,
800
- "acceptance_method": "typical_acceptance_sampler",
801
- },
802
- "enable_chunked_prefill": False
803
- }
804
- # Try a range of common k.
805
- for k in [1, 2, 3]
806
- ] + [{
807
- "speculative_config": {
808
- "model": "JackFram/llama-68m",
809
- "num_speculative_tokens": k,
810
- "acceptance_method": "typical_acceptance_sampler",
811
- },
812
- "enable_chunked_prefill": True,
813
- "max_num_batched_tokens": 4,
814
- "max_num_seqs": 4
815
- } for k in [1, 2, 3]])
816
-@pytest.mark.parametrize("batch_size", [1, 32])
817
-@pytest.mark.parametrize(
818
- "output_len",
819
- [
820
- # Use smaller output len for fast test.
821
- 32,
822
- ])
823
-@pytest.mark.parametrize("seed", [1])
824
-@create_new_process_for_each_test()
825
-def test_typical_acceptance_sampling(vllm_runner, common_llm_kwargs,
826
- per_test_common_llm_kwargs,
827
- baseline_llm_kwargs, test_llm_kwargs,
828
- batch_size: int, output_len: int,
829
- seed: int):
830
- """Verify that speculative decoding produces exact equality to without spec
831
- decode with TypicalAcceptanceSampler as the draft token acceptance
832
- sampling method.
833
- """
834
- run_equality_correctness_test(vllm_runner,
835
- common_llm_kwargs,
836
- per_test_common_llm_kwargs,
837
- baseline_llm_kwargs,
838
- test_llm_kwargs,
839
- batch_size,
840
- max_output_len=output_len,
841
- seed=seed,
842
- temperature=0.0)
tests/spec_decode/e2e/test_ngram_correctness.py
deleted
-392
@@ -1,392 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""This docstring details important information on the testing methodology.
4
-
5
-Most of the tests rely on "greedy equality", where we expect the output of
6
-speculative decoding on a sequence to exactly match the output of normal non-
7
-speculative decoding.
8
-
9
-Since speculative decoding with rejection sampling guarantees that the output
10
-distribution matches the target model's output distribution (up to hardware
11
-numerics, see https://arxiv.org/pdf/2302.01318.pdf), we can expect greedy
12
-equality.
13
-
14
-For ngram lookup, its idea comes from https://github.com/apoorvumang/prompt-lookup-decoding,
15
-and is merged into transform code base: https://github.com/huggingface/transformers/pull/27775.
16
-Since there is no model is needed for generate the proposal, we could make
17
-the testcase much simpler than drafter multi-step one.
18
-
19
-However, we still need to verify below scenario could be passed:
20
- * Batch size 1 greedy equality
21
- * Batch size >1 greedy equality
22
- * Test greedy equality under preemption
23
- * Test greedy equality under various ngram sizes / speculative sizes
24
-
25
-With those tests, we can say at least, ngram spec would not break the
26
-correctness for the target model outputs.
27
-"""
28
-
29
-import pytest
30
-
31
-from ..utils import maybe_enable_chunked_prefill
32
-from .conftest import run_equality_correctness_test
33
-
34
-
35
-@pytest.mark.parametrize(
36
- "common_llm_kwargs",
37
- [{
38
- # Skip cuda graph recording for fast test.
39
- "enforce_eager": True,
40
-
41
- # Print spec metrics.
42
- "disable_log_stats": False,
43
-
44
- # The original model is float32, keep it for numerical stability.
45
- "dtype": "float32",
46
- }])
47
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [
48
- {
49
- "model_name": "JackFram/llama-68m",
50
- },
51
-])
52
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
53
-@pytest.mark.parametrize("test_llm_kwargs", [
54
- {
55
- "speculative_config": {
56
- "method": "ngram",
57
- "num_speculative_tokens": 5,
58
- "prompt_lookup_max": 3,
59
- "disable_mqa_scorer": False,
60
- },
61
- },
62
- {
63
- "speculative_config": {
64
- "method": "ngram",
65
- "num_speculative_tokens": 5,
66
- "prompt_lookup_max": 3,
67
- "disable_mqa_scorer": True,
68
- },
69
- },
70
-])
71
-@pytest.mark.parametrize("output_len", [
72
- 256,
73
-])
74
-@pytest.mark.parametrize("batch_size", [1, 32])
75
-@pytest.mark.parametrize("prefill_chunk_size", [-1, 4])
76
-@pytest.mark.parametrize("seed", [1])
77
-def test_ngram_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
78
- per_test_common_llm_kwargs,
79
- baseline_llm_kwargs, test_llm_kwargs,
80
- batch_size: int, output_len: int,
81
- prefill_chunk_size: int, seed: int):
82
- """Verify greedy equality on a tiny model with different batch size."""
83
- maybe_enable_chunked_prefill(prefill_chunk_size, common_llm_kwargs)
84
- run_equality_correctness_test(vllm_runner,
85
- common_llm_kwargs,
86
- per_test_common_llm_kwargs,
87
- baseline_llm_kwargs,
88
- test_llm_kwargs,
89
- batch_size,
90
- max_output_len=output_len,
91
- seed=seed,
92
- temperature=0.0)
93
-
94
-
95
-@pytest.mark.parametrize(
96
- "common_llm_kwargs",
97
- [{
98
- # Skip cuda graph recording for fast test.
99
- "enforce_eager": True,
100
-
101
- # Print spec metrics.
102
- "disable_log_stats": False,
103
-
104
- # The original model is float32, keep it for numerical stability.
105
- "dtype": "float32",
106
- }])
107
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [
108
- {
109
- "model_name": "JackFram/llama-68m",
110
- },
111
-])
112
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
113
-@pytest.mark.parametrize("test_llm_kwargs", [
114
- {
115
- "speculative_config": {
116
- "method": "ngram",
117
- "num_speculative_tokens": 5,
118
- "prompt_lookup_max": 3,
119
- "disable_logprobs": False,
120
- },
121
- },
122
- {
123
- "speculative_config": {
124
- "method": "ngram",
125
- "num_speculative_tokens": 5,
126
- "prompt_lookup_max": 3,
127
- "disable_logprobs": True,
128
- },
129
- },
130
-])
131
-@pytest.mark.parametrize("output_len", [
132
- 8,
133
-])
134
-@pytest.mark.parametrize("batch_size", [8])
135
-@pytest.mark.parametrize("seed", [1])
136
-@pytest.mark.parametrize("logprobs", [1, 6])
137
-def test_ngram_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
138
- per_test_common_llm_kwargs,
139
- baseline_llm_kwargs, test_llm_kwargs,
140
- batch_size: int, output_len: int, seed: int,
141
- logprobs: int):
142
- """Verify greedy equality on a tiny model with different batch size."""
143
- run_equality_correctness_test(
144
- vllm_runner,
145
- common_llm_kwargs,
146
- per_test_common_llm_kwargs,
147
- baseline_llm_kwargs,
148
- test_llm_kwargs,
149
- batch_size,
150
- max_output_len=output_len,
151
- seed=seed,
152
- temperature=0.0,
153
- logprobs=logprobs,
154
- prompt_logprobs=logprobs,
155
- disable_logprobs=test_llm_kwargs["speculative_config"]
156
- ["disable_logprobs"])
157
-
158
-
159
-@pytest.mark.parametrize(
160
- "common_llm_kwargs",
161
- [{
162
- "block_size": 16,
163
- # 2 for small prompt, 256//8 for generated.
164
- "num_gpu_blocks_override": 2 + 256 // 8,
165
- "max_model_len": (2 + 256 // 8) * 8,
166
-
167
- # Skip cuda graph recording for fast test.
168
- "enforce_eager": True,
169
-
170
- # The original model is float32, keep it for numerical stability.
171
- "dtype": "float32",
172
- }])
173
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [
174
- {
175
- "model_name": "JackFram/llama-160m",
176
- },
177
-])
178
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
179
-@pytest.mark.parametrize("test_llm_kwargs", [
180
- {
181
- "speculative_config": {
182
- "method": "ngram",
183
- "num_speculative_tokens": 5,
184
- "prompt_lookup_max": 3,
185
- },
186
- "enable_chunked_prefill": False,
187
- },
188
- {
189
- "speculative_config": {
190
- "method": "ngram",
191
- "num_speculative_tokens": 5,
192
- "prompt_lookup_max": 3,
193
- "disable_mqa_scorer": True,
194
- },
195
- "enable_chunked_prefill": True,
196
- "max_num_batched_tokens": 4,
197
- "max_num_seqs": 4
198
- },
199
-])
200
-@pytest.mark.parametrize(
201
- "output_len",
202
- [
203
- # Use small output len for fast test.
204
- 256,
205
- ])
206
-@pytest.mark.parametrize("batch_size", [4])
207
-@pytest.mark.parametrize("seed", [1])
208
-def test_ngram_e2e_greedy_correctness_with_preemption(
209
- vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
210
- baseline_llm_kwargs, test_llm_kwargs, batch_size: int, output_len: int,
211
- seed: int):
212
- """Verify greedy equality, even when some sequences are preempted mid-
213
- generation.
214
- """
215
- run_equality_correctness_test(vllm_runner,
216
- common_llm_kwargs,
217
- per_test_common_llm_kwargs,
218
- baseline_llm_kwargs,
219
- test_llm_kwargs,
220
- batch_size,
221
- max_output_len=output_len,
222
- temperature=0,
223
- seed=seed)
224
-
225
-
226
-@pytest.mark.parametrize(
227
- "common_llm_kwargs",
228
- [{
229
- "model_name": "JackFram/llama-68m",
230
-
231
- # Skip cuda graph recording for fast test.
232
- "enforce_eager": True,
233
-
234
- # The original model is float32, keep it for numerical stability.
235
- "dtype": "float32",
236
- }])
237
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
238
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
239
-@pytest.mark.parametrize(
240
- "test_llm_kwargs",
241
- [
242
- {
243
- "speculative_config": {
244
- "method": "ngram",
245
- "num_speculative_tokens": k,
246
- "prompt_lookup_max": 3,
247
- },
248
- }
249
- # Try a range of common k, as well as large speculation.
250
- for k in [1, 3, 5]
251
- ] + [
252
- {
253
- "speculative_config": {
254
- "method": "ngram",
255
- "num_speculative_tokens": k,
256
- "prompt_lookup_max": 1,
257
- },
258
- }
259
- # Try a range of common k, as well as large speculation.
260
- for k in [1, 3, 5]
261
- ])
262
-@pytest.mark.parametrize("batch_size", [2])
263
-@pytest.mark.parametrize(
264
- "output_len",
265
- [
266
- # Use smaller output len for fast test.
267
- 32,
268
- ])
269
-@pytest.mark.parametrize("seed", [1])
270
-def test_ngram_different_k(vllm_runner, common_llm_kwargs,
271
- per_test_common_llm_kwargs, baseline_llm_kwargs,
272
- test_llm_kwargs, batch_size: int, output_len: int,
273
- seed: int):
274
- """Verify that ngram speculative decoding produces exact equality
275
- to without spec decode with many different values of k and
276
- different ngram prompt_lookup_max.
277
- """
278
- run_equality_correctness_test(vllm_runner,
279
- common_llm_kwargs,
280
- per_test_common_llm_kwargs,
281
- baseline_llm_kwargs,
282
- test_llm_kwargs,
283
- batch_size,
284
- max_output_len=output_len,
285
- seed=seed,
286
- temperature=0.0)
287
-
288
-
289
-@pytest.mark.parametrize(
290
- "common_llm_kwargs",
291
- [{
292
- "model_name": "JackFram/llama-68m",
293
-
294
- # Skip cuda graph recording for fast test.
295
- "enforce_eager": True,
296
-
297
- # The original model is float32, keep it for numerical stability.
298
- "dtype": "float32",
299
- }])
300
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
301
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
302
-@pytest.mark.parametrize("test_llm_kwargs", [{
303
- "speculative_config": {
304
- "method": "ngram",
305
- "num_speculative_tokens": 5,
306
- "prompt_lookup_max": 3,
307
- "disable_by_batch_size": 4
308
- },
309
-}, {
310
- "speculative_config": {
311
- "method": "ngram",
312
- "num_speculative_tokens": 5,
313
- "prompt_lookup_max": 3,
314
- "disable_by_batch_size": 4,
315
- "disable_mqa_scorer": True,
316
- },
317
- "enable_chunked_prefill": True,
318
- "max_num_batched_tokens": 4,
319
- "max_num_seqs": 4
320
-}])
321
-@pytest.mark.parametrize("batch_size", [1, 5])
322
-@pytest.mark.parametrize(
323
- "output_len",
324
- [
325
- # Use smaller output len for fast test.
326
- 32,
327
- ])
328
-@pytest.mark.parametrize("seed", [1])
329
-def test_ngram_disable_queue(vllm_runner, common_llm_kwargs,
330
- per_test_common_llm_kwargs, baseline_llm_kwargs,
331
- test_llm_kwargs, batch_size: int, output_len: int,
332
- seed: int):
333
- """Verify that ngram speculative decoding produces exact equality
334
- to without spec decode with many different values of k and
335
- different ngram prompt_lookup_max.
336
- """
337
- run_equality_correctness_test(vllm_runner,
338
- common_llm_kwargs,
339
- per_test_common_llm_kwargs,
340
- baseline_llm_kwargs,
341
- test_llm_kwargs,
342
- batch_size,
343
- max_output_len=output_len,
344
- seed=seed,
345
- temperature=0.0)
346
-
347
-
348
-@pytest.mark.parametrize(
349
- "common_llm_kwargs",
350
- [{
351
- "model_name": "JackFram/llama-68m",
352
-
353
- # Skip cuda graph recording for fast test.
354
- "enforce_eager": True,
355
-
356
- # The original model is float32, keep it for numerical stability.
357
- "dtype": "float32",
358
- }])
359
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
360
-@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
361
-@pytest.mark.parametrize("test_llm_kwargs", [{
362
- "speculative_config": {
363
- "method": "ngram",
364
- "num_speculative_tokens": 5,
365
- "prompt_lookup_max": 3,
366
- "disable_mqa_scorer": True,
367
- },
368
-}])
369
-@pytest.mark.parametrize("batch_size", [1, 5])
370
-@pytest.mark.parametrize(
371
- "output_len",
372
- [
373
- # Use smaller output len for fast test.
374
- 32,
375
- ])
376
-@pytest.mark.parametrize("seed", [1])
377
-def test_ngram_scorer(vllm_runner, common_llm_kwargs,
378
- per_test_common_llm_kwargs, baseline_llm_kwargs,
379
- test_llm_kwargs, batch_size: int, output_len: int,
380
- seed: int):
381
- """Verify that ngram speculative decoding generates the same output
382
- with batch expansion scorer and mqa scorer.
383
- """
384
- run_equality_correctness_test(vllm_runner,
385
- common_llm_kwargs,
386
- per_test_common_llm_kwargs,
387
- baseline_llm_kwargs,
388
- test_llm_kwargs,
389
- batch_size,
390
- max_output_len=output_len,
391
- seed=seed,
392
- temperature=0.0)
tests/spec_decode/e2e/test_seed.py
deleted
-70
@@ -1,70 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import pytest
5
-
6
-from .conftest import run_equality_correctness_test
7
-
8
-# main model
9
-MAIN_MODEL = "JackFram/llama-68m"
10
-
11
-# speculative model
12
-SPEC_MODEL = "JackFram/llama-160m"
13
-
14
-
15
-@pytest.mark.parametrize(
16
- "common_llm_kwargs",
17
- [{
18
- "model_name": "JackFram/llama-68m",
19
-
20
- # Skip cuda graph recording for fast test.
21
- "enforce_eager": True,
22
-
23
- # speculative config
24
- "speculative_config": {
25
- "model": "JackFram/llama-160m",
26
- "num_speculative_tokens": 3,
27
- },
28
- }])
29
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
30
-@pytest.mark.parametrize("baseline_llm_kwargs", [{"seed": 1}])
31
-@pytest.mark.parametrize("test_llm_kwargs", [{"seed": 5}])
32
-@pytest.mark.parametrize("batch_size", [1, 8, 32])
33
-@pytest.mark.parametrize("temperature", [0.1, 1.0])
34
-@pytest.mark.parametrize(
35
- "output_len",
36
- [
37
- # Use smaller output len for fast test.
38
- 20,
39
- ])
40
-def test_seeded_consistency(vllm_runner, common_llm_kwargs,
41
- per_test_common_llm_kwargs, baseline_llm_kwargs,
42
- test_llm_kwargs, batch_size: int,
43
- temperature: float, output_len: int):
44
- """Verify outputs are consistent across multiple runs with same seed
45
- """
46
- run_equality_correctness_test(
47
- vllm_runner,
48
- common_llm_kwargs,
49
- per_test_common_llm_kwargs,
50
- baseline_llm_kwargs,
51
- test_llm_kwargs,
52
- batch_size,
53
- max_output_len=output_len,
54
- temperature=temperature,
55
- disable_seed=False,
56
- )
57
-
58
- # Ensure this same test does fail if we _don't_ include per-request seeds
59
- with pytest.raises(AssertionError):
60
- run_equality_correctness_test(
61
- vllm_runner,
62
- common_llm_kwargs,
63
- per_test_common_llm_kwargs,
64
- baseline_llm_kwargs,
65
- test_llm_kwargs,
66
- batch_size,
67
- max_output_len=output_len,
68
- temperature=temperature,
69
- disable_seed=True,
70
- )
tests/spec_decode/test_batch_expansion.py
deleted
-110
@@ -1,110 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import pytest
5
-import torch
6
-
7
-from vllm.spec_decode.batch_expansion import BatchExpansionTop1Scorer
8
-
9
-from .utils import create_seq_group_metadata_from_prompts, mock_worker
10
-
11
-
12
-@pytest.mark.parametrize('num_target_seq_ids', [100])
13
-@pytest.mark.skip_global_cleanup
14
-def test_create_target_seq_id_iterator(num_target_seq_ids: int):
15
- """Verify all new sequence ids are greater than all input
16
- seq ids.
17
- """
18
- scorer = BatchExpansionTop1Scorer(mock_worker(), 'cuda:0', 32_000)
19
-
20
- all_seq_ids = [
21
- [1, 3, 5, 7],
22
- list(range(100)) + [0],
23
- [100],
24
- ]
25
-
26
- for seq_ids in all_seq_ids:
27
- max_seq_id = max(seq_ids)
28
- iterator = scorer._create_target_seq_id_iterator(seq_ids) # pylint: disable=protected-access
29
- for _ in range(num_target_seq_ids):
30
- assert next(iterator) > max_seq_id
31
-
32
-
33
-@pytest.mark.parametrize('k', [1, 2, 6])
34
-@pytest.mark.skip_global_cleanup
35
-def test_get_token_ids_to_score(k: int):
36
- """Verify correct tokens are selected for scoring.
37
- """
38
- proposal_token_ids = torch.tensor(
39
- list(range(k)),
40
- dtype=torch.int64,
41
- device='cuda',
42
- )
43
-
44
- expected_output: list[list[int]] = [
45
- [],
46
- ]
47
- for i in range(proposal_token_ids.shape[0]):
48
- expected_output.append(proposal_token_ids[:i + 1].tolist())
49
-
50
- scorer = BatchExpansionTop1Scorer(mock_worker(), 'cuda:0', 32_000)
51
- actual_output = scorer._get_token_ids_to_score(proposal_token_ids.tolist()) # pylint: disable=protected-access
52
-
53
- actual_output = [
54
- x.tolist() if isinstance(x, torch.Tensor) else x for x in actual_output
55
- ]
56
-
57
- assert actual_output == expected_output
58
-
59
-
60
-@pytest.mark.parametrize('k', [1, 2, 6])
61
-@pytest.mark.skip_global_cleanup
62
-def test_create_single_target_seq_group_metadata(k: int):
63
- """Verify correct creation of a batch-expanded seq group metadata.
64
- """
65
-
66
- prompt_tokens = [1, 2, 3]
67
- prev_output_tokens = [4, 5, 6]
68
-
69
- token_ids = list(range(k))
70
-
71
- num_tokens_processed = len(prompt_tokens) + len(prev_output_tokens) - 1
72
-
73
- final_seq_len = len(prompt_tokens) + len(prev_output_tokens) + len(
74
- token_ids)
75
-
76
- block_size = 32
77
- input_seq_group_metadata = create_seq_group_metadata_from_prompts(
78
- [prompt_tokens], 2048 // block_size, block_size, [final_seq_len],
79
- [prev_output_tokens], [num_tokens_processed])[0]
80
-
81
- input_seq_id = list(input_seq_group_metadata.seq_data.keys())[0]
82
- target_seq_id = 100
83
-
84
- scorer = BatchExpansionTop1Scorer(mock_worker(), 'cuda:0', 32_000)
85
- output = scorer._create_single_target_seq_group_metadata( # pylint: disable=protected-access
86
- input_seq_group_metadata,
87
- input_seq_id,
88
- target_seq_id,
89
- token_ids,
90
- input_seq_group_metadata.sampling_params,
91
- )
92
-
93
- assert output.request_id == input_seq_group_metadata.request_id
94
- assert output.sampling_params.repetition_penalty == \
95
- input_seq_group_metadata.sampling_params.repetition_penalty
96
- assert output.sampling_params.temperature == \
97
- input_seq_group_metadata.sampling_params.temperature
98
- assert output.sampling_params.top_p == \
99
- input_seq_group_metadata.sampling_params.top_p
100
- assert output.sampling_params.top_k == \
101
- input_seq_group_metadata.sampling_params.top_k
102
- assert len(output.seq_data) == 1
103
- assert output.seq_data[target_seq_id].get_prompt_token_ids() == tuple(
104
- prompt_tokens)
105
- assert output.seq_data[target_seq_id].get_output_token_ids() == tuple(
106
- prev_output_tokens + token_ids)
107
-
108
- assert len(output.block_tables) == 1
109
- assert output.block_tables[
110
- target_seq_id] == input_seq_group_metadata.block_tables[input_seq_id]
tests/spec_decode/test_dynamic_spec_decode.py
deleted
-90
@@ -1,90 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from unittest.mock import MagicMock, patch
5
-
6
-import pytest
7
-import torch
8
-
9
-from vllm.sequence import ExecuteModelRequest
10
-from vllm.spec_decode.metrics import AsyncMetricsCollector
11
-from vllm.spec_decode.multi_step_worker import MultiStepWorker
12
-from vllm.spec_decode.spec_decode_worker import SpecDecodeWorker
13
-from vllm.spec_decode.top1_proposer import Top1Proposer
14
-
15
-from .test_utils import mock_spec_decode_sampler
16
-from .utils import create_batch, mock_worker
17
-
18
-
19
-@pytest.mark.parametrize('queue_size', [4])
20
-@pytest.mark.parametrize('batch_size', [1])
21
-@pytest.mark.parametrize('k', [1])
22
-@pytest.mark.parametrize("acceptance_sampler_method",
23
- ["rejection_sampler", "typical_acceptance_sampler"])
24
-@torch.inference_mode()
25
-def test_disable_spec_tokens(queue_size: int, batch_size: int, k: int,
26
- acceptance_sampler_method: str):
27
- """Verify that speculative tokens are disabled when the batch size
28
- exceeds the threshold.
29
- """
30
- disable_by_batch_size = 3
31
- draft_worker = mock_worker(cls=MultiStepWorker)
32
- target_worker = mock_worker()
33
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
34
- worker = SpecDecodeWorker(proposer_worker=draft_worker,
35
- scorer_worker=target_worker,
36
- spec_decode_sampler=mock_spec_decode_sampler(
37
- acceptance_sampler_method),
38
- disable_logprobs=False,
39
- metrics_collector=metrics_collector,
40
- disable_by_batch_size=disable_by_batch_size)
41
-
42
- exception_secret = 'artificial stop'
43
- draft_worker.get_spec_proposals.side_effect = ValueError(exception_secret)
44
-
45
- seq_group_metadata_list, _, _ = create_batch(batch_size, k)
46
- execute_model_req = ExecuteModelRequest(
47
- seq_group_metadata_list=seq_group_metadata_list,
48
- num_lookahead_slots=k,
49
- running_queue_size=queue_size)
50
-
51
- if queue_size > disable_by_batch_size:
52
- with patch.object(worker,
53
- '_run_no_spec',
54
- side_effect=ValueError(exception_secret)), \
55
- pytest.raises(ValueError, match=exception_secret):
56
- worker.execute_model(execute_model_req=execute_model_req)
57
-
58
- # When the batch size is larger than the threshold,
59
- # we expect no speculative tokens (0).
60
- expected_num_spec_tokens = None if queue_size < disable_by_batch_size else 0
61
- assert seq_group_metadata_list[
62
- 0].num_speculative_tokens == expected_num_spec_tokens
63
-
64
- draft_worker.sampler_output.side_effect = ValueError(exception_secret)
65
-
66
- proposer = Top1Proposer(
67
- worker=draft_worker,
68
- device='cpu', # not used
69
- vocab_size=100, # not used
70
- # Must be long enough to avoid being skipped due to length.
71
- max_proposal_len=1024,
72
- )
73
-
74
- if queue_size < disable_by_batch_size:
75
- # Should raise exception when executing the mocked draft model.
76
- with pytest.raises(ValueError, match=exception_secret):
77
- proposer.get_spec_proposals(
78
- execute_model_req=ExecuteModelRequest(
79
- seq_group_metadata_list=seq_group_metadata_list,
80
- num_lookahead_slots=k),
81
- seq_ids_with_bonus_token_in_last_step=set())
82
- else:
83
- # Should not execute the draft model because spec decode is disabled
84
- # for all requests. Accordingly, the proposal length should be 0.
85
- proposals = proposer.get_spec_proposals(
86
- execute_model_req=ExecuteModelRequest(
87
- seq_group_metadata_list=seq_group_metadata_list,
88
- num_lookahead_slots=k),
89
- seq_ids_with_bonus_token_in_last_step=set())
90
- assert proposals.proposal_lens.tolist() == [0] * batch_size
tests/spec_decode/test_memory_usage.py
deleted
-91
@@ -1,91 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""This docstring details important information on the testing methodology.
4
-
5
-This test verifies that memory usage remains constant (or never grows) when
6
-we enable / disable speculation via --speculative-disable-by-batch-size.
7
-
8
-There are a lot of things we try to keep track of between batches of requests
9
-and if certain tensors are not freed from memory, can result in CUDA ooms.
10
-
11
-This is particularly relevant for production situations where speculation might
12
-be enabled during off hours, but disabled once traffic peaks during the workday.
13
-Since traffic will stay high for a long period of time, verifying we do not
14
-increase our memory usage over time is essential to prevent possible CUDA ooms.
15
-"""
16
-
17
-import torch
18
-
19
-import vllm
20
-from tests.core.utils import create_dummy_prompt
21
-from vllm.sequence import SequenceGroup
22
-
23
-ITERATIONS = 100
24
-MAIN_MODEL = "JackFram/llama-68m"
25
-
26
-# speculative model
27
-SPEC_MODEL = "abhigoyal/vllm-medusa-llama-68m-random"
28
-
29
-BATCH_SIZE = 5
30
-SPEC_DISABLE_BATCH_SIZE = 2
31
-
32
-
33
-def add_seq_group_to_engine(engine: vllm.LLMEngine, seq_group: SequenceGroup):
34
- scheduler = engine.scheduler[0]
35
- scheduler.add_seq_group(seq_group)
36
-
37
-
38
-"""
39
-Since we are using a batch size greater than the disabled batch size,
40
-we can ensure we go through the _no_spec codepath for most of our engine steps.
41
-"""
42
-
43
-
44
-def test_memory_usage_no_spec():
45
- previous_memory_allocated = None
46
- llm = vllm.LLM(model=MAIN_MODEL,
47
- speculative_config={
48
- "model": SPEC_MODEL,
49
- "num_speculative_tokens": 3,
50
- "disable_by_batch_size": SPEC_DISABLE_BATCH_SIZE,
51
- })
52
-
53
- batch_sequences = set()
54
- engine = llm.llm_engine
55
-
56
- for i in range(ITERATIONS):
57
- seq, seq_group = create_dummy_prompt(request_id=str(i),
58
- prompt_length=10,
59
- min_tokens=10,
60
- max_tokens=10)
61
-
62
- add_seq_group_to_engine(engine, seq_group)
63
-
64
- batch_sequences.add(seq)
65
- engine.step()
66
- for seq in list(batch_sequences):
67
- if seq.is_finished():
68
- batch_sequences.remove(seq)
69
-
70
- # If we aren't at our batch size yet, continue
71
- if len(batch_sequences) <= BATCH_SIZE:
72
- continue
73
-
74
- # Otherwise, loop until at least one request is done
75
- while not any(seq.is_finished() for seq in batch_sequences):
76
- engine.step()
77
-
78
- # Remove it from the set
79
- for seq in list(batch_sequences):
80
- if seq.is_finished():
81
- batch_sequences.remove(seq)
82
-
83
- # At this point, we are always at the case where we have finished
84
- # processing some number of requests from the batch after running
85
- # several _no_spec executions. The memory should not have
86
- # increased between the previous time this was recorded and the
87
- # current time.
88
- if previous_memory_allocated is None:
89
- previous_memory_allocated = torch.cuda.memory_allocated()
90
- else:
91
- assert previous_memory_allocated == torch.cuda.memory_allocated()
tests/spec_decode/test_metrics.py
deleted
-205
@@ -1,205 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import math
5
-from unittest.mock import MagicMock
6
-
7
-import pytest
8
-import torch
9
-
10
-from vllm.spec_decode.metrics import AsyncMetricsCollector
11
-
12
-
13
-def test_initial_call_returns_none():
14
- """Expect first call to get metrics to return None.
15
- """
16
- spec_decode_sampler = MagicMock()
17
- spec_decode_sampler.num_accepted_tokens = torch.tensor(0,
18
- dtype=torch.long,
19
- device='cuda')
20
- spec_decode_sampler.num_emitted_tokens = torch.tensor(0,
21
- dtype=torch.long,
22
- device='cuda')
23
- spec_decode_sampler.num_draft_tokens = 0
24
-
25
- collector = AsyncMetricsCollector(spec_decode_sampler)
26
- collector.init_gpu_tensors(rank=0)
27
- maybe_metrics = collector.maybe_collect_rejsample_metrics(k=5)
28
- assert maybe_metrics is None
29
-
30
-
31
-def test_second_call_returns_metrics():
32
- """Expect second call to not return None.
33
- """
34
- spec_decode_sampler = MagicMock()
35
- spec_decode_sampler.num_accepted_tokens = torch.tensor(0,
36
- dtype=torch.long,
37
- device='cuda')
38
- spec_decode_sampler.num_emitted_tokens = torch.tensor(0,
39
- dtype=torch.long,
40
- device='cuda')
41
- spec_decode_sampler.num_draft_tokens = 0
42
-
43
- collect_interval_s = 5.0
44
- timer = MagicMock()
45
- timer.side_effect = [
46
- 0.0, collect_interval_s + 0.1, collect_interval_s + 0.2
47
- ]
48
-
49
- collector = AsyncMetricsCollector(spec_decode_sampler=spec_decode_sampler,
50
- timer=timer,
51
- collect_interval_s=collect_interval_s)
52
- collector.init_gpu_tensors(rank=0)
53
- _ = collector.maybe_collect_rejsample_metrics(k=5)
54
- metrics = collector.maybe_collect_rejsample_metrics(k=5)
55
- assert metrics is not None
56
-
57
-
58
-@pytest.mark.parametrize("rank", [1, 2, 3, 4])
59
-def test_nonzero_rank_noop(rank):
60
- """Verify nonzero ranks don't collect metrics.
61
- """
62
- spec_decode_sampler = MagicMock()
63
- spec_decode_sampler.num_accepted_tokens = torch.tensor(0,
64
- dtype=torch.long,
65
- device='cuda')
66
- spec_decode_sampler.num_emitted_tokens = torch.tensor(0,
67
- dtype=torch.long,
68
- device='cuda')
69
- spec_decode_sampler.num_draft_tokens = 0
70
-
71
- collector = AsyncMetricsCollector(spec_decode_sampler)
72
- collector.init_gpu_tensors(rank=rank)
73
- _ = collector.maybe_collect_rejsample_metrics(k=5)
74
- metrics = collector.maybe_collect_rejsample_metrics(k=5)
75
- assert metrics is None
76
-
77
-
78
-def test_noop_until_time():
79
- """Verify metrics aren't collected until enough time passes.
80
- """
81
- spec_decode_sampler = MagicMock()
82
- spec_decode_sampler.num_accepted_tokens = torch.tensor(0,
83
- dtype=torch.long,
84
- device='cuda')
85
- spec_decode_sampler.num_emitted_tokens = torch.tensor(0,
86
- dtype=torch.long,
87
- device='cuda')
88
- spec_decode_sampler.num_draft_tokens = 0
89
-
90
- collect_interval_s = 5.0
91
- timer = MagicMock()
92
- timer.side_effect = [
93
- 0.0, collect_interval_s - 0.1, collect_interval_s - 0.1,
94
- collect_interval_s + 0.1, collect_interval_s + 0.1
95
- ]
96
-
97
- collector = AsyncMetricsCollector(spec_decode_sampler=spec_decode_sampler,
98
- timer=timer,
99
- collect_interval_s=collect_interval_s)
100
- collector.init_gpu_tensors(rank=0)
101
-
102
- _ = collector.maybe_collect_rejsample_metrics(k=5)
103
- metrics = collector.maybe_collect_rejsample_metrics(k=5)
104
- assert metrics is None
105
-
106
- _ = collector.maybe_collect_rejsample_metrics(k=5)
107
- metrics = collector.maybe_collect_rejsample_metrics(k=5)
108
- assert metrics is not None
109
-
110
-
111
-def test_timer_is_reset():
112
- """Verify that the internal timer inside AsyncMetricsCollector
113
- is reset after collection.
114
- """
115
- spec_decode_sampler = MagicMock()
116
- spec_decode_sampler.num_accepted_tokens = torch.tensor(0,
117
- dtype=torch.long,
118
- device='cuda')
119
- spec_decode_sampler.num_emitted_tokens = torch.tensor(0,
120
- dtype=torch.long,
121
- device='cuda')
122
- spec_decode_sampler.num_draft_tokens = 0
123
-
124
- collect_interval_s = 5.0
125
- timer = MagicMock()
126
- timer.side_effect = [
127
- 0.0,
128
- collect_interval_s + 0.1,
129
- collect_interval_s + 0.1,
130
- collect_interval_s + 0.2,
131
- collect_interval_s + 0.2,
132
- 2 * collect_interval_s + 0.1,
133
- 2 * collect_interval_s + 0.1,
134
- ]
135
-
136
- collector = AsyncMetricsCollector(spec_decode_sampler=spec_decode_sampler,
137
- timer=timer,
138
- collect_interval_s=collect_interval_s)
139
- collector.init_gpu_tensors(rank=0)
140
-
141
- _ = collector.maybe_collect_rejsample_metrics(k=5)
142
- metrics = collector.maybe_collect_rejsample_metrics(k=5)
143
- assert metrics is not None
144
-
145
- _ = collector.maybe_collect_rejsample_metrics(k=5)
146
- metrics = collector.maybe_collect_rejsample_metrics(k=5)
147
- assert metrics is None
148
-
149
- _ = collector.maybe_collect_rejsample_metrics(k=5)
150
- metrics = collector.maybe_collect_rejsample_metrics(k=5)
151
- assert metrics is not None
152
-
153
-
154
-@pytest.mark.parametrize("has_data", [True, False])
155
-def test_initial_metrics_has_correct_values(has_data: bool):
156
- """Test correctness of metrics data.
157
- """
158
- if has_data:
159
- num_accepted_tokens = 103
160
- num_emitted_tokens = 104
161
- num_draft_tokens = 105
162
- else:
163
- num_accepted_tokens = 0
164
- num_emitted_tokens = 0
165
- num_draft_tokens = 0
166
- k = 5
167
-
168
- max_num_emitted_tokens = AsyncMetricsCollector.get_max_num_emitted_tokens(
169
- num_draft_tokens, k)
170
-
171
- spec_decode_sampler = MagicMock()
172
- spec_decode_sampler.num_accepted_tokens = torch.tensor(num_accepted_tokens,
173
- dtype=torch.long,
174
- device='cuda')
175
- spec_decode_sampler.num_emitted_tokens = torch.tensor(num_emitted_tokens,
176
- dtype=torch.long,
177
- device='cuda')
178
- spec_decode_sampler.num_draft_tokens = num_draft_tokens
179
-
180
- collect_interval_s = 5.0
181
- timer = MagicMock()
182
- timer.side_effect = [
183
- 0.0, collect_interval_s + 0.1, collect_interval_s + 0.2
184
- ]
185
-
186
- collector = AsyncMetricsCollector(spec_decode_sampler=spec_decode_sampler,
187
- timer=timer,
188
- collect_interval_s=collect_interval_s)
189
- collector.init_gpu_tensors(rank=0)
190
- _ = collector.maybe_collect_rejsample_metrics(k)
191
- metrics = collector.maybe_collect_rejsample_metrics(k)
192
-
193
- assert metrics.num_spec_tokens == k
194
- assert metrics.accepted_tokens == num_accepted_tokens
195
- assert metrics.draft_tokens == num_draft_tokens
196
- assert metrics.emitted_tokens == num_emitted_tokens
197
-
198
- if has_data:
199
- assert (metrics.draft_acceptance_rate == num_accepted_tokens /
200
- num_draft_tokens)
201
- assert (metrics.system_efficiency == num_emitted_tokens /
202
- max_num_emitted_tokens)
203
- else:
204
- assert math.isnan(metrics.draft_acceptance_rate)
205
- assert math.isnan(metrics.system_efficiency)
tests/spec_decode/test_multi_step_worker.py
deleted
-838
@@ -1,838 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import random
5
-from unittest.mock import MagicMock
6
-
7
-import pytest
8
-import torch
9
-
10
-from vllm.attention.selector import (_Backend,
11
- global_force_attn_backend_context_manager)
12
-from vllm.model_executor.layers.sampler import SamplerOutput
13
-from vllm.model_executor.utils import set_random_seed
14
-from vllm.sequence import (ExecuteModelRequest, HiddenStates, Logprob,
15
- get_all_seq_ids)
16
-from vllm.spec_decode.draft_model_runner import TP1DraftModelRunner
17
-from vllm.spec_decode.multi_step_worker import MultiStepWorker
18
-from vllm.spec_decode.top1_proposer import Top1Proposer
19
-from vllm.worker.worker import Worker
20
-
21
-from .utils import (assert_logprobs_dict_allclose, create_batch,
22
- create_seq_group_metadata_from_prompts, create_worker,
23
- patch_execute_model_with_seeds, zero_kv_cache)
24
-
25
-
26
-@pytest.mark.parametrize('num_steps', list(range(1, 17)))
27
-def test_assert_enough_kv_space(num_steps: int):
28
- """Test that the multi step worker checks for sufficient space in the KV
29
- cache. It should throw if it cannot run all the steps.
30
- """
31
- block_size = 16
32
- num_gpu_blocks = 2048 // block_size
33
-
34
- prompts = [
35
- list(range(block_size * 3)),
36
- list(range(block_size * 2)),
37
- ]
38
-
39
- prev_output_tokens = [
40
- list(range(block_size * 1)),
41
- list(range(block_size * 2)),
42
- ]
43
-
44
- final_prompt_lens = [
45
- len(prompt + output) + num_steps
46
- for prompt, output in zip(prompts, prev_output_tokens)
47
- ]
48
-
49
- inputs = create_seq_group_metadata_from_prompts(
50
- prompts,
51
- num_gpu_blocks,
52
- block_size,
53
- final_prompt_lens,
54
- continuations=prev_output_tokens)
55
-
56
- assert_enough_kv_space = MultiStepWorker._assert_enough_kv_space # pylint: disable=protected-access
57
- worker = MagicMock()
58
- worker.model_runner.block_size = block_size
59
-
60
- for seq_group_metadata in inputs:
61
- original_block_tables = seq_group_metadata.block_tables
62
-
63
- # No exception.
64
- assert_enough_kv_space(worker, inputs, num_steps)
65
-
66
- seq_group_metadata.block_tables = {
67
- seq_id: []
68
- for seq_id, physical_blocks in original_block_tables.items()
69
- }
70
-
71
- # Expect exception.
72
- with pytest.raises(ValueError,
73
- match='times but found insufficient KV space for'):
74
- assert_enough_kv_space(worker, inputs, num_steps)
75
-
76
- seq_group_metadata.block_tables = original_block_tables
77
-
78
-
79
-@torch.inference_mode()
80
-def test_same_output_for_single_step():
81
- """Verify the multi step worker produces the same output as the normal
82
- worker for num_steps=1.
83
- """
84
- seed = 100
85
- model_name = 'JackFram/llama-68m'
86
-
87
- block_size = 32
88
- num_gpu_blocks = 2048 // block_size
89
- multi_step_worker = create_worker(
90
- MultiStepWorker,
91
- model_name,
92
- block_size,
93
- num_gpu_blocks,
94
- seed,
95
- model_runner_cls=TP1DraftModelRunner,
96
- )
97
- worker = create_worker(
98
- Worker,
99
- model_name,
100
- block_size,
101
- num_gpu_blocks,
102
- seed,
103
- )
104
- # multi_step_worker.model_runner = worker.model_runner
105
- # multi_step_worker.cache_engine = worker.cache_engine
106
-
107
- num_steps = 1
108
-
109
- prompts = [
110
- [1, 2, 3, 4, 5],
111
- [6, 7, 8, 9, 10],
112
- ]
113
-
114
- final_prompt_lens = [len(prompt) + num_steps for prompt in prompts]
115
-
116
- multi_step_seq_group = create_seq_group_metadata_from_prompts(
117
- prompts,
118
- num_gpu_blocks,
119
- block_size,
120
- final_prompt_lens=final_prompt_lens)
121
-
122
- zero_kv_cache(multi_step_worker.cache_engine)
123
- set_random_seed(seed)
124
- actual_output, _ = multi_step_worker.sampler_output(
125
- execute_model_req=ExecuteModelRequest(
126
- seq_group_metadata_list=multi_step_seq_group),
127
- sample_len=num_steps,
128
- seq_ids_with_bonus_token_in_last_step=set())
129
- assert len(actual_output) == num_steps
130
- actual_output = actual_output[0]
131
-
132
- single_step_seq_group = create_seq_group_metadata_from_prompts(
133
- prompts,
134
- num_gpu_blocks,
135
- block_size,
136
- final_prompt_lens=final_prompt_lens)
137
-
138
- zero_kv_cache(worker.cache_engine)
139
- set_random_seed(seed)
140
- expected_output = worker.execute_model(
141
- execute_model_req=ExecuteModelRequest(
142
- seq_group_metadata_list=single_step_seq_group))[0]
143
-
144
- actual_token_ids = [
145
- output.samples[0].output_token for output in actual_output
146
- ]
147
- actual_logprobs = [output.samples[0].logprobs for output in actual_output]
148
-
149
- expected_token_ids = [
150
- output.samples[0].output_token for output in expected_output
151
- ]
152
- expected_logprobs = [
153
- output.samples[0].logprobs for output in expected_output
154
- ]
155
-
156
- assert actual_token_ids == expected_token_ids
157
-
158
- print(f'{actual_logprobs=}')
159
- print(f'{expected_logprobs=}')
160
- assert_logprobs_dict_allclose(actual_logprobs, expected_logprobs)
161
-
162
-
163
-@torch.inference_mode()
164
-def test_same_output_for_multi_step():
165
- """Verify the multi-step worker produces the same output as the normal
166
- worker when num_steps > 1. This test runs the multi-step worker once, and
167
- then runs the worker num_steps times, and compares the output.
168
- """
169
- seed = 100
170
- model_name = 'JackFram/llama-68m'
171
-
172
- block_size = 16
173
- num_gpu_blocks = 2048 // block_size
174
- multi_step_worker = create_worker(
175
- MultiStepWorker,
176
- model_name,
177
- block_size,
178
- num_gpu_blocks,
179
- seed,
180
- )
181
-
182
- worker = create_worker(
183
- Worker,
184
- model_name,
185
- block_size,
186
- num_gpu_blocks,
187
- seed,
188
- )
189
-
190
- # Make sure we go over the block boundary.
191
- num_steps = block_size + 1
192
-
193
- random.seed(seed)
194
- prompts = [[
195
- random.randint(0, 1000) for _ in range(random.randint(10, 20))
196
- ] for _ in range(10)]
197
-
198
- final_prompt_lens = [len(prompt) + num_steps for prompt in prompts]
199
-
200
- rand_seeds = list(random.randint(0, 100) for _ in range(num_steps))
201
- multi_step_worker.execute_model = patch_execute_model_with_seeds(
202
- multi_step_worker, rand_seeds)
203
- worker.execute_model = patch_execute_model_with_seeds(worker, rand_seeds)
204
-
205
- continuations = [[1] for _ in prompts]
206
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
207
- prompts,
208
- num_gpu_blocks,
209
- block_size,
210
- continuations=continuations,
211
- final_prompt_lens=final_prompt_lens)
212
-
213
- # Run multi-step.
214
- zero_kv_cache(multi_step_worker.cache_engine)
215
- set_random_seed(seed)
216
- multi_step_output, _ = multi_step_worker.sampler_output(
217
- execute_model_req=ExecuteModelRequest(
218
- seq_group_metadata_list=seq_group_metadata_list),
219
- sample_len=num_steps,
220
- seq_ids_with_bonus_token_in_last_step=set())
221
-
222
- # Run single-step repeatedly.
223
- zero_kv_cache(worker.cache_engine)
224
- single_step_output: list[SamplerOutput] = []
225
- continuations = [[1] for _ in prompts]
226
- set_random_seed(seed)
227
-
228
- for _ in multi_step_output:
229
-
230
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
231
- prompts,
232
- num_gpu_blocks,
233
- block_size,
234
- continuations=continuations,
235
- final_prompt_lens=final_prompt_lens)
236
-
237
- single_step_output.extend(
238
- worker.execute_model(execute_model_req=ExecuteModelRequest(
239
- seq_group_metadata_list=seq_group_metadata_list)))
240
-
241
- # Append output tokens to new sequence data.
242
- for i, seq_group_output in enumerate(single_step_output[-1]):
243
- continuations[i].append(seq_group_output.samples[0].output_token)
244
-
245
- # Get token ids and logprobs for comparison.
246
- multi_step_output_logprobs: list[list[dict[int,
247
- Logprob]]] = [[]
248
- for _ in prompts]
249
- single_step_output_logprobs: list[list[dict[int,
250
- Logprob]]] = [[]
251
- for _ in prompts]
252
-
253
- multi_step_output_token_ids: list[list[int]] = [[] for _ in prompts]
254
- single_step_output_token_ids: list[list[int]] = [[] for _ in prompts]
255
- for i, _ in enumerate(prompts):
256
- for multi_step, single_step in zip(multi_step_output,
257
- single_step_output):
258
- multi_step_output_token_ids[i].append(
259
- multi_step[i].samples[0].output_token)
260
- single_step_output_token_ids[i].append(
261
- single_step[i].samples[0].output_token)
262
-
263
- multi_step_output_logprobs[i].append(
264
- multi_step[i].samples[0].logprobs)
265
- single_step_output_logprobs[i].append(
266
- single_step[i].samples[0].logprobs)
267
-
268
- # Print per-sequence token ids
269
- for i, (multi_step_tokens, single_step_tokens) in enumerate(
270
- zip(multi_step_output_token_ids, single_step_output_token_ids)):
271
- print(f'{i=} {multi_step_tokens=}')
272
- print(f'{i=} {single_step_tokens=}')
273
- print(f'{i=} equal {multi_step_tokens == single_step_tokens}')
274
-
275
- # Assert token ids are equal.
276
- for multi_step_tokens, single_step_tokens in zip(
277
- multi_step_output_token_ids, single_step_output_token_ids):
278
- assert multi_step_tokens == single_step_tokens
279
-
280
- # Assert logprobs are equal.
281
- for multi_step_logprobs, single_step_logprobs in zip(
282
- multi_step_output_logprobs, single_step_output_logprobs):
283
- assert_logprobs_dict_allclose(multi_step_logprobs,
284
- single_step_logprobs)
285
-
286
-
287
-@torch.inference_mode()
288
-def test_multi_step_with_batch_expansion_correct_output():
289
- """
290
- In this test we verify that the MultiStepWorker is able to handle bonus
291
- tokens correctly. The test verifies that if a sequence has a
292
- bonus token then the MultiStepWorker is able to expand the batch by adding
293
- new sequences corresponding to the sequences with bonus tokens. The
294
- expanded batch is then used for predicting the next tokens.
295
- """
296
- seed = 100
297
- model_name = 'JackFram/llama-68m'
298
-
299
- block_size = 16
300
- num_gpu_blocks = 2048 // block_size
301
- batch_size = 128
302
- multi_step_worker = create_worker(
303
- MultiStepWorker,
304
- model_name,
305
- block_size,
306
- num_gpu_blocks,
307
- seed,
308
- model_runner_cls=TP1DraftModelRunner,
309
- )
310
- multi_step_worker.set_include_gpu_probs_tensor()
311
- worker = create_worker(
312
- Worker,
313
- model_name,
314
- block_size,
315
- num_gpu_blocks,
316
- seed,
317
- )
318
- random.seed(seed)
319
- prompts = [[0] for _ in range(batch_size)]
320
- num_steps = 2
321
- final_prompt_lens = [(num_steps + 1) for prompt in prompts]
322
- rand_seeds = list(random.randint(0, 100) for _ in range(num_steps))
323
- multi_step_worker.execute_model = patch_execute_model_with_seeds(
324
- multi_step_worker, rand_seeds)
325
- worker.execute_model = patch_execute_model_with_seeds(worker, rand_seeds)
326
- # Create the test continuations
327
- continuations = [[random.randint(0, 1000)] for _ in prompts]
328
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
329
- prompts,
330
- num_gpu_blocks,
331
- block_size,
332
- continuations=continuations,
333
- final_prompt_lens=final_prompt_lens)
334
-
335
- # Run single-step twice to generate 2 tokens. This
336
- # will simulate the bonus token case with the second token
337
- # being the bonus token.
338
- zero_kv_cache(worker.cache_engine)
339
- single_step_output: list[SamplerOutput] = []
340
- set_random_seed(seed)
341
- for _ in range(num_steps):
342
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
343
- prompts,
344
- num_gpu_blocks,
345
- block_size,
346
- continuations=continuations,
347
- final_prompt_lens=final_prompt_lens)
348
- single_step_output.extend(
349
- worker.execute_model(execute_model_req=ExecuteModelRequest(
350
- seq_group_metadata_list=seq_group_metadata_list)))
351
- # Append output tokens to new sequence data.
352
- for i, seq_group_output in enumerate(single_step_output[-1]):
353
- continuations[i].append(seq_group_output.samples[0].output_token)
354
-
355
- # Create continuations for the MultiStepWorker. The continuations have
356
- # 2 tokens in order to simulate the bonus token case.
357
- multi_step_continuations = []
358
- for continuation in continuations:
359
- multi_step_continuations.append(continuation[:2])
360
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
361
- prompts,
362
- num_gpu_blocks,
363
- block_size,
364
- continuations=multi_step_continuations,
365
- final_prompt_lens=final_prompt_lens)
366
-
367
- # Run multi-step and verify that the third token prediction is accurate
368
- # for all sequences.
369
- zero_kv_cache(multi_step_worker.cache_engine)
370
- all_seq_ids = {i for i in range(batch_size)}
371
- multi_step_output, _ = multi_step_worker.sampler_output(
372
- execute_model_req=ExecuteModelRequest(
373
- seq_group_metadata_list=seq_group_metadata_list),
374
- sample_len=1,
375
- seq_ids_with_bonus_token_in_last_step=all_seq_ids)
376
- for index, output in enumerate(multi_step_output[-1].outputs):
377
- assert (continuations[index][-1] == output.samples[0].output_token)
378
-
379
-
380
-@torch.inference_mode()
381
-def test_multi_step_with_batch_expansion_incorrect_output():
382
- """
383
- Tests the MultiStepWorker's ability to handle batch expansion with bonus
384
- tokens in a negative case scenario. This test provides the MultiStepWorker
385
- with a batch containing sequences with bonus tokens but specifies the
386
- sequence IDs with bonus tokens incorrectly. The test verifies that the
387
- MultiStepWorker generates correct tokens for the sequences where the
388
- sequence ID is specified correctly and incorrect tokens for those where
389
- the sequence ID is specified incorrectly.
390
- """
391
- seed = 100
392
- model_name = 'JackFram/llama-68m'
393
-
394
- block_size = 16
395
- num_gpu_blocks = 2048 // block_size
396
- batch_size = 128
397
- multi_step_worker = create_worker(
398
- MultiStepWorker,
399
- model_name,
400
- block_size,
401
- num_gpu_blocks,
402
- seed,
403
- model_runner_cls=TP1DraftModelRunner,
404
- )
405
- multi_step_worker.set_include_gpu_probs_tensor()
406
- worker = create_worker(
407
- Worker,
408
- model_name,
409
- block_size,
410
- num_gpu_blocks,
411
- seed,
412
- )
413
- random.seed(seed)
414
- prompts = [[0] for _ in range(batch_size)]
415
- num_steps = 2
416
- final_prompt_lens = [(num_steps + 1) for prompt in prompts]
417
- rand_seeds = list(random.randint(0, 100) for _ in range(num_steps))
418
- multi_step_worker.execute_model = patch_execute_model_with_seeds(
419
- multi_step_worker, rand_seeds)
420
- worker.execute_model = patch_execute_model_with_seeds(worker, rand_seeds)
421
- # Create the test continuations
422
- continuations = [[random.randint(0, 1000)] for _ in prompts]
423
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
424
- prompts,
425
- num_gpu_blocks,
426
- block_size,
427
- continuations=continuations,
428
- final_prompt_lens=final_prompt_lens)
429
- # Run single-step twice to generate 2 tokens. This
430
- # will simulate the bonus token case with the second token
431
- # being the bonus token.
432
- zero_kv_cache(worker.cache_engine)
433
- single_step_output: list[SamplerOutput] = []
434
- set_random_seed(seed)
435
- for _ in range(num_steps):
436
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
437
- prompts,
438
- num_gpu_blocks,
439
- block_size,
440
- continuations=continuations,
441
- final_prompt_lens=final_prompt_lens)
442
- single_step_output.extend(
443
- worker.execute_model(execute_model_req=ExecuteModelRequest(
444
- seq_group_metadata_list=seq_group_metadata_list)))
445
- # Append output tokens to new sequence data.
446
- for i, seq_group_output in enumerate(single_step_output[-1]):
447
- continuations[i].append(seq_group_output.samples[0].output_token)
448
-
449
- # Create continuations for the MultiStepWorker. The continuations have
450
- # 2 tokens in order to simulate the bonus token case.
451
- multi_step_continuations = []
452
- for continuation in continuations:
453
- multi_step_continuations.append(continuation[:2])
454
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
455
- prompts,
456
- num_gpu_blocks,
457
- block_size,
458
- continuations=multi_step_continuations,
459
- final_prompt_lens=final_prompt_lens)
460
-
461
- # Run multi-step. In this run INCORRECTLY specify that only the odd number
462
- # sequences have bonus tokens. Verify that with this setting the third token
463
- # prediction is accurate only for the odd numbered sequences. Also verify
464
- # that the prediction might be wrong for some of the even numbered
465
- # sequences.
466
- zero_kv_cache(multi_step_worker.cache_engine)
467
- set_random_seed(seed)
468
- odd_seq_ids = {i for i in range(batch_size) if i % 2 != 0}
469
- multi_step_output, _ = multi_step_worker.sampler_output(
470
- execute_model_req=ExecuteModelRequest(
471
- seq_group_metadata_list=seq_group_metadata_list),
472
- sample_len=1,
473
- seq_ids_with_bonus_token_in_last_step=odd_seq_ids)
474
- num_mismatch = 0
475
- for index, output in enumerate(multi_step_output[-1].outputs):
476
- if (index % 2) != 0:
477
- assert (continuations[index][-1] == output.samples[0].output_token)
478
- elif (continuations[index][-1] != output.samples[0].output_token):
479
- num_mismatch += 1
480
- # The prediction is accurate for some of the sequences even without proper
481
- # handling of the bonus tokens. Hence verify that the number of sequences
482
- # for which there is a mismatch is > 0.
483
- assert (num_mismatch > 0)
484
-
485
-
486
-@torch.inference_mode()
487
-@pytest.mark.parametrize('num_steps', [1, 2, 3, 4])
488
-# The choice of backends forces the multi_step_worker to choose between
489
-# the vanilla model_runner and TP1DraftModelRunner and that we can test
490
-# both code paths.
491
-@pytest.mark.parametrize('attn_backend',
492
- [_Backend.XFORMERS, _Backend.FLASH_ATTN])
493
-def test_multi_step_correct_kvcache(num_steps, attn_backend):
494
- """Verify that the KV cache of the draft model
495
- is correctly updated for sequences with bonus token.
496
- """
497
- seed = 100
498
- model_name = "JackFram/llama-68m"
499
-
500
- block_size = 16
501
- num_gpu_blocks = 2048 // block_size
502
- batch_size = 1
503
-
504
- with global_force_attn_backend_context_manager(attn_backend):
505
- dtype = 'float16' if attn_backend == _Backend.FLASH_ATTN else 'float32'
506
- multi_step_worker = create_worker(MultiStepWorker,
507
- model_name,
508
- block_size,
509
- num_gpu_blocks,
510
- seed,
511
- model_runner_cls=TP1DraftModelRunner,
512
- dtype=dtype)
513
- multi_step_worker.set_include_gpu_probs_tensor()
514
- worker = create_worker(Worker,
515
- model_name,
516
- block_size,
517
- num_gpu_blocks,
518
- seed,
519
- dtype=dtype)
520
-
521
- prompts = [[0] for _ in range(batch_size)]
522
- # Already generate two tokens for the sequence
523
- # so that we can simulate the bonus token case
524
- multi_step_continuations = [[
525
- random.randint(0, 1000),
526
- random.randint(0, 1000)
527
- ] for _ in prompts]
528
- final_prompt_lens = [len(prompt) + 2 + num_steps for prompt in prompts]
529
-
530
- seq_ids_with_bonus_token_in_last_step = set(range(batch_size))
531
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
532
- prompts,
533
- num_gpu_blocks,
534
- block_size,
535
- continuations=multi_step_continuations,
536
- final_prompt_lens=final_prompt_lens)
537
-
538
- # Run multi-step.
539
- zero_kv_cache(multi_step_worker.cache_engine)
540
- multi_step_worker.sampler_output(execute_model_req=ExecuteModelRequest(
541
- seq_group_metadata_list=seq_group_metadata_list),
542
- sample_len=num_steps,
543
- seq_ids_with_bonus_token_in_last_step=
544
- seq_ids_with_bonus_token_in_last_step)
545
-
546
- # Run single-step repeatedly.
547
- zero_kv_cache(worker.cache_engine)
548
- # Generate the kv cache for the bonus token first
549
- single_step_continuations = [c[:1] for c in multi_step_continuations]
550
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
551
- prompts,
552
- num_gpu_blocks,
553
- block_size,
554
- continuations=single_step_continuations,
555
- final_prompt_lens=final_prompt_lens)
556
- single_step_output = worker.execute_model(
557
- execute_model_req=ExecuteModelRequest(
558
- seq_group_metadata_list=seq_group_metadata_list))
559
- for _ in range(num_steps):
560
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
561
- prompts,
562
- num_gpu_blocks,
563
- block_size,
564
- continuations=multi_step_continuations,
565
- final_prompt_lens=final_prompt_lens)
566
-
567
- single_step_output = worker.execute_model(
568
- execute_model_req=ExecuteModelRequest(
569
- seq_group_metadata_list=seq_group_metadata_list))
570
-
571
- for i, seq_group_output in enumerate(single_step_output[-1]):
572
- multi_step_continuations[i].append(
573
- seq_group_output.samples[0].output_token)
574
-
575
- # Verify that the KV cache of the single-step and
576
- # multi-step workers are the same.
577
- single_step_gpu_cache = worker.cache_engine[0].gpu_cache
578
- multi_step_gpu_cache = multi_step_worker.cache_engine[0].gpu_cache
579
- num_layers = len(single_step_gpu_cache)
580
- allclose = lambda a, b: torch.allclose(
581
- a.cuda(), b.cuda(), rtol=1e-2, atol=1e-2)
582
- for i in range(num_layers):
583
- assert allclose(single_step_gpu_cache[i][0],
584
- multi_step_gpu_cache[i][0])
585
- assert allclose(single_step_gpu_cache[i][1],
586
- multi_step_gpu_cache[i][1])
587
-
588
-
589
-@torch.inference_mode()
590
-def test_draft_proposals_full_speculation_len():
591
- """Verify Top1Proposer correctly handles case where all sequences
592
- can speculate.
593
- """
594
- k = 10
595
- batch_size = 32
596
- vocab_size = 32_000
597
- device = 'cuda:0'
598
-
599
- draft_worker = MagicMock()
600
- proposer = Top1Proposer(
601
- worker=draft_worker,
602
- device=device,
603
- vocab_size=vocab_size,
604
- max_proposal_len=2048,
605
- )
606
- draft_worker.sampler_output.return_value = [
607
- SamplerOutput(
608
- outputs=[],
609
- sampled_token_probs=torch.rand(batch_size,
610
- vocab_size,
611
- device=device,
612
- dtype=torch.float32),
613
- logprobs=torch.rand(batch_size,
614
- vocab_size,
615
- device=device,
616
- dtype=torch.float32),
617
- sampled_token_ids=torch.randint(low=0,
618
- high=vocab_size,
619
- size=(batch_size, ),
620
- device=device,
621
- dtype=torch.long),
622
- ) for _ in range(k)
623
- ], True
624
-
625
- seq_group_metadata_list, _, _ = create_batch(batch_size, k)
626
-
627
- proposals = proposer.get_spec_proposals(
628
- execute_model_req=ExecuteModelRequest(
629
- seq_group_metadata_list=seq_group_metadata_list,
630
- num_lookahead_slots=k),
631
- seq_ids_with_bonus_token_in_last_step=set())
632
-
633
- assert torch.is_tensor(proposals.proposal_token_ids)
634
- assert torch.is_tensor(proposals.proposal_probs)
635
-
636
- assert proposals.proposal_token_ids.shape == torch.Size([batch_size, k])
637
- assert proposals.proposal_probs.shape[:-1] == torch.Size([batch_size, k])
638
-
639
- assert proposals.proposal_lens.shape == torch.Size([batch_size])
640
- assert proposals.proposal_lens.tolist() == [k for _ in range(batch_size)]
641
-
642
-
643
-@torch.inference_mode()
644
-def test_draft_proposals_no_speculations():
645
- """Verify Top1Proposer correctly handles case where no sequences
646
- can speculate.
647
- """
648
- k = 10
649
- batch_size = 32
650
- vocab_size = 32_000
651
- device = 'cuda:0'
652
- prompt_len = 10
653
-
654
- draft_worker = MagicMock()
655
- proposer = Top1Proposer(
656
- worker=draft_worker,
657
- device=device,
658
- vocab_size=vocab_size,
659
- max_proposal_len=prompt_len + k - 1,
660
- )
661
-
662
- seq_group_metadata_list, _, _ = create_batch(batch_size,
663
- k,
664
- prompt_len=prompt_len)
665
-
666
- proposals = proposer.get_spec_proposals(
667
- execute_model_req=ExecuteModelRequest(
668
- seq_group_metadata_list=seq_group_metadata_list,
669
- num_lookahead_slots=k),
670
- seq_ids_with_bonus_token_in_last_step=set())
671
-
672
- assert torch.is_tensor(proposals.proposal_token_ids)
673
- assert torch.is_tensor(proposals.proposal_probs)
674
-
675
- assert proposals.proposal_token_ids.shape == torch.Size([batch_size, k])
676
- assert proposals.proposal_probs.shape[:-1] == torch.Size([batch_size, k])
677
-
678
- assert proposals.proposal_lens.shape == torch.Size([batch_size])
679
- assert proposals.proposal_lens.tolist() == [0 for _ in range(batch_size)]
680
-
681
-
682
-@torch.inference_mode()
683
-def test_draft_proposals_mixed_k():
684
- """Verify Top1Proposer correctly handles case some sequences can
685
- speculate and some can't.
686
- """
687
- k = 10
688
- batch_size = 32
689
- vocab_size = 32_000
690
- device = 'cuda:0'
691
-
692
- small_prompt_len = 5
693
- long_prompt_len = 10
694
- prev_output_token_len = 20
695
-
696
- expected_num_proposal_seqs = 6
697
- expected_num_no_proposal_seqs = batch_size - expected_num_proposal_seqs
698
-
699
- prompt_len = [
700
- small_prompt_len for _ in range(expected_num_proposal_seqs - 1)
701
- ] + [long_prompt_len
702
- for _ in range(expected_num_no_proposal_seqs)] + [small_prompt_len]
703
-
704
- draft_worker = MagicMock()
705
- proposer = Top1Proposer(
706
- worker=draft_worker,
707
- device=device,
708
- vocab_size=vocab_size,
709
- max_proposal_len=long_prompt_len + prev_output_token_len + k - 1,
710
- )
711
-
712
- draft_worker.sampler_output.return_value = [
713
- SamplerOutput(
714
- outputs=[],
715
- sampled_token_probs=torch.rand(expected_num_proposal_seqs,
716
- vocab_size,
717
- device=device,
718
- dtype=torch.float32),
719
- logprobs=torch.rand(expected_num_proposal_seqs,
720
- vocab_size,
721
- device=device,
722
- dtype=torch.float32),
723
- sampled_token_ids=torch.randint(
724
- low=0,
725
- high=vocab_size,
726
- size=(expected_num_proposal_seqs, ),
727
- device=device,
728
- dtype=torch.long),
729
- ) for _ in range(k)
730
- ], True
731
-
732
- seq_group_metadata_list, _, _ = create_batch(
733
- batch_size,
734
- k,
735
- prompt_len=prompt_len,
736
- prev_output_token_len=prev_output_token_len,
737
- )
738
-
739
- proposals = proposer.get_spec_proposals(
740
- execute_model_req=ExecuteModelRequest(
741
- seq_group_metadata_list=seq_group_metadata_list,
742
- num_lookahead_slots=k),
743
- seq_ids_with_bonus_token_in_last_step=set())
744
-
745
- assert torch.is_tensor(proposals.proposal_token_ids)
746
- assert torch.is_tensor(proposals.proposal_probs)
747
-
748
- assert proposals.proposal_token_ids.shape == torch.Size([batch_size, k])
749
- assert proposals.proposal_probs.shape[:-1] == torch.Size([batch_size, k])
750
-
751
- assert proposals.proposal_lens.shape == torch.Size([batch_size])
752
- assert proposals.proposal_lens.tolist() == [
753
- k for _ in range(expected_num_proposal_seqs - 1)
754
- ] + [0 for _ in range(expected_num_no_proposal_seqs)] + [k]
755
-
756
-
757
-@torch.inference_mode()
758
-def test_use_draft_model_runner_advance_step():
759
- """Verify that draft model runner triggers advance step
760
- when applicable.
761
- """
762
- seed = 100
763
- model_name = 'JackFram/llama-68m'
764
-
765
- k = 5
766
- batch_size = 32
767
- block_size = 32
768
- num_gpu_blocks = 2048 // block_size
769
- worker = create_worker(
770
- MultiStepWorker,
771
- model_name,
772
- block_size,
773
- num_gpu_blocks,
774
- seed,
775
- model_runner_cls=TP1DraftModelRunner,
776
- )
777
-
778
- # Mock "_gpu_advance_step" to raise an exception when called.
779
- exception_secret = "artificial stop"
780
- worker.model_runner._gpu_advance_step = MagicMock()
781
- worker.model_runner._gpu_advance_step.side_effect = ValueError(
782
- exception_secret)
783
-
784
- seq_group_metadata_list, _, _ = create_batch(batch_size,
785
- k,
786
- block_size=block_size,
787
- num_gpu_blocks=num_gpu_blocks)
788
-
789
- # Fallback (should not call) when num_steps=1.
790
- execute_model_req = ExecuteModelRequest(
791
- seq_group_metadata_list=seq_group_metadata_list,
792
- num_lookahead_slots=k,
793
- num_steps=1)
794
- worker.execute_model(execute_model_req=execute_model_req)
795
-
796
- # Expect exception if _gpu_advance_step is called.
797
- execute_model_req = ExecuteModelRequest(
798
- seq_group_metadata_list=seq_group_metadata_list,
799
- num_lookahead_slots=k,
800
- num_steps=k)
801
-
802
- with pytest.raises(ValueError, match=exception_secret):
803
- worker.execute_model(execute_model_req=execute_model_req)
804
- call_args_list = worker.model_runner._gpu_advance_step.call_args_list
805
- assert len(call_args_list) == 1
806
-
807
-
808
-@torch.inference_mode()
809
-def test_expand_execute_model_request_sync_with_expand_hidden_states():
810
- """
811
- In this test we verify that the logic for expanding the
812
- seq_group_metadata_list remains in sync with the expansion logic of
813
- the HiddenStates in _expand_execute_model_request.
814
- """
815
- k = 5
816
- batch_size = 16
817
- seq_with_bonus_token_in_last_step = [1, 3, 8, 10, 13, 15]
818
-
819
- seq_group_metadata_list, _, _ = create_batch(batch_size, k)
820
-
821
- execute_model_request = ExecuteModelRequest(
822
- seq_group_metadata_list,
823
- previous_hidden_states=HiddenStates(
824
- torch.arange(batch_size), seq_group_metadata_list,
825
- torch.arange(batch_size, 2 * batch_size)))
826
-
827
- expanded_execute_model_request, orig_seq_group_ids = MultiStepWorker.\
828
- _expand_execute_model_request(execute_model_request,
829
- seq_with_bonus_token_in_last_step)
830
-
831
- all_seq_ids = torch.tensor(
832
- get_all_seq_ids(
833
- expanded_execute_model_request.seq_group_metadata_list))
834
- ref_expanded_hidden_states = all_seq_ids + batch_size
835
- ref_expanded_hidden_states[orig_seq_group_ids] -= batch_size
836
-
837
- assert (ref_expanded_hidden_states == expanded_execute_model_request.
838
- previous_hidden_states.hidden_states).all().item()
tests/spec_decode/test_ngram_worker.py
deleted
-221
@@ -1,221 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import torch
5
-
6
-from vllm.sequence import ExecuteModelRequest
7
-from vllm.spec_decode.ngram_worker import NGramWorker
8
-from vllm.spec_decode.top1_proposer import Top1Proposer
9
-
10
-from .utils import create_seq_group_metadata_from_prompts, create_worker
11
-
12
-
13
-def test_ngram_algo_correctness_for_single_no_match():
14
- """Verify our ngram algo find the right candidate in the prompt
15
-
16
- For the scenario cannot find any candidate in one single batch
17
- """
18
- block_size = 32
19
- num_gpu_blocks = 2048 // block_size
20
- seed = 100
21
- model_name = 'JackFram/llama-68m'
22
- vocab_size = 32_000
23
- device = 'cuda:0'
24
-
25
- ngram_worker = create_worker(
26
- NGramWorker,
27
- model_name,
28
- block_size,
29
- num_gpu_blocks,
30
- seed,
31
- )
32
-
33
- proposer = Top1Proposer(
34
- worker=ngram_worker,
35
- device=device,
36
- vocab_size=vocab_size,
37
- max_proposal_len=20,
38
- )
39
-
40
- # set ngram window [1, 3], which is window=1/2/3
41
- ngram_worker.set_ngram_window_size(1, 3)
42
-
43
- prompts = [
44
- # shall find no candidate
45
- [1, 2, 3, 4, 5, 6, 7],
46
- ]
47
-
48
- proposal_len = 5
49
- final_prompt_lens = [len(prompt) + proposal_len for prompt in prompts]
50
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
51
- prompts,
52
- num_gpu_blocks,
53
- block_size,
54
- final_prompt_lens=final_prompt_lens)
55
-
56
- proposals = proposer.get_spec_proposals(
57
- execute_model_req=ExecuteModelRequest(
58
- seq_group_metadata_list=seq_group_metadata_list,
59
- num_lookahead_slots=proposal_len),
60
- seq_ids_with_bonus_token_in_last_step=None)
61
-
62
- assert torch.is_tensor(proposals.proposal_token_ids)
63
- assert torch.is_tensor(proposals.proposal_probs)
64
-
65
- assert proposals.proposal_token_ids.shape == torch.Size([1, proposal_len])
66
- assert proposals.proposal_probs.shape[:-1] == torch.Size([1, proposal_len])
67
- assert proposals.proposal_lens.shape == torch.Size([1])
68
- assert proposals.proposal_lens.tolist() == [0]
69
-
70
-
71
-def test_ngram_algo_correctness_for_batches_not_match_all():
72
- """Verify our ngram algo find the right candidate in the prompt
73
-
74
- For the scenario find some candidate not full in batchs
75
- """
76
- block_size = 32
77
- num_gpu_blocks = 2048 // block_size
78
- seed = 100
79
- model_name = 'JackFram/llama-68m'
80
- vocab_size = 32_000
81
- device = 'cuda:0'
82
-
83
- ngram_worker = create_worker(
84
- NGramWorker,
85
- model_name,
86
- block_size,
87
- num_gpu_blocks,
88
- seed,
89
- )
90
-
91
- proposer = Top1Proposer(
92
- worker=ngram_worker,
93
- device=device,
94
- vocab_size=vocab_size,
95
- max_proposal_len=20,
96
- )
97
-
98
- # set ngram window [1, 3], which is window=1/2/3
99
- ngram_worker.set_ngram_window_size(1, 3)
100
-
101
- prompts = [
102
- # shall find no candidate
103
- [1, 2, 3, 4, 5, 6, 7],
104
- # shall find candidate 12,13,14,15,16
105
- [11, 12, 13, 14, 15, 16, 11],
106
- # shall find candidate 23,24,25,26,21
107
- [21, 21, 22, 23, 24, 25, 26, 21, 22],
108
- # shall find candidate 34,35,36,37,38
109
- [31, 32, 31, 32, 33, 34, 35, 36, 37, 38, 31, 32, 33],
110
- # shall find no candidate as exceed max_proposal_len
111
- [
112
- 31, 32, 31, 32, 31, 32, 31, 32, 31, 32, 31, 32, 33, 34, 35, 36, 37,
113
- 38, 31, 32, 33
114
- ],
115
- ]
116
-
117
- proposal_len = 5
118
- final_prompt_lens = [len(prompt) + proposal_len for prompt in prompts]
119
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
120
- prompts,
121
- num_gpu_blocks,
122
- block_size,
123
- final_prompt_lens=final_prompt_lens)
124
- for sg in seq_group_metadata_list:
125
- sg.is_prompt = False
126
- proposals = proposer.get_spec_proposals(
127
- execute_model_req=ExecuteModelRequest(
128
- seq_group_metadata_list=seq_group_metadata_list,
129
- num_lookahead_slots=proposal_len),
130
- seq_ids_with_bonus_token_in_last_step=None)
131
-
132
- assert torch.is_tensor(proposals.proposal_token_ids)
133
- assert torch.is_tensor(proposals.proposal_probs)
134
-
135
- assert proposals.proposal_token_ids.shape == torch.Size([5, proposal_len])
136
- assert proposals.proposal_probs.shape[:-1] == torch.Size([5, proposal_len])
137
- assert proposals.proposal_lens.shape == torch.Size([5])
138
-
139
- # the first sequence has no match so proposal_len should be overwritten to 0
140
- assert proposals.proposal_lens.tolist(
141
- ) == [0] + [proposal_len for _ in range(3)] + [0]
142
-
143
- for i in range(proposal_len):
144
- assert proposals.proposal_token_ids[0][i] == -1
145
- assert proposals.proposal_token_ids[1][i] == prompts[1][i + 1]
146
- assert proposals.proposal_token_ids[2][i] == prompts[2][i + 3]
147
- assert proposals.proposal_token_ids[3][i] == prompts[3][i + 5]
148
- assert proposals.proposal_token_ids[4][i] == -1
149
-
150
-
151
-def test_ngram_algo_correctness_for_batches_match_all():
152
- """Verify our ngram algo find the right candidate in the prompt
153
-
154
- For the scenario find candidate in all batches
155
- """
156
-
157
- block_size = 32
158
- num_gpu_blocks = 2048 // block_size
159
- seed = 100
160
- model_name = 'JackFram/llama-68m'
161
- vocab_size = 32_000
162
- device = 'cuda:0'
163
-
164
- ngram_worker = create_worker(
165
- NGramWorker,
166
- model_name,
167
- block_size,
168
- num_gpu_blocks,
169
- seed,
170
- )
171
-
172
- proposer = Top1Proposer(
173
- worker=ngram_worker,
174
- device=device,
175
- vocab_size=vocab_size,
176
- max_proposal_len=20,
177
- )
178
-
179
- # set ngram window [0, 3], which is window=1/2/3
180
- ngram_worker.set_ngram_window_size(1, 3)
181
-
182
- prompts = [
183
- # shall find candidate 12,13,14,15,16
184
- [11, 12, 13, 14, 15, 16, 11],
185
- # shall find candidate 23,24,25,26,21
186
- [21, 21, 22, 23, 24, 25, 26, 21, 22],
187
- # shall find candidate 34,35,36,37,38
188
- [31, 32, 31, 32, 33, 34, 35, 36, 37, 38, 31, 32, 33],
189
- ]
190
-
191
- proposal_len = 5
192
- final_prompt_lens = [len(prompt) + proposal_len for prompt in prompts]
193
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
194
- prompts,
195
- num_gpu_blocks,
196
- block_size,
197
- final_prompt_lens=final_prompt_lens)
198
-
199
- # Normally drafter is run on decode requests only; here we check the output
200
- # of the ngram worker as it is the sole proposer that has no forward.
201
- for sg in seq_group_metadata_list:
202
- sg.is_prompt = False
203
- proposals = proposer.get_spec_proposals(
204
- execute_model_req=ExecuteModelRequest(
205
- seq_group_metadata_list=seq_group_metadata_list,
206
- num_lookahead_slots=proposal_len),
207
- seq_ids_with_bonus_token_in_last_step=None)
208
-
209
- assert torch.is_tensor(proposals.proposal_token_ids)
210
- assert torch.is_tensor(proposals.proposal_probs)
211
-
212
- assert proposals.proposal_token_ids.shape == torch.Size([3, proposal_len])
213
- assert proposals.proposal_probs.shape[:-1] == torch.Size([3, proposal_len])
214
- assert proposals.proposal_lens.shape == torch.Size([3])
215
-
216
- assert proposals.proposal_lens.tolist() == [proposal_len for _ in range(3)]
217
-
218
- for i in range(proposal_len):
219
- assert proposals.proposal_token_ids[0][i] == prompts[0][i + 1]
220
- assert proposals.proposal_token_ids[1][i] == prompts[1][i + 3]
221
- assert proposals.proposal_token_ids[2][i] == prompts[2][i + 5]
tests/spec_decode/test_scorer.py
deleted
-116
@@ -1,116 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import random
5
-
6
-import pytest
7
-import torch
8
-
9
-from vllm.sequence import ExecuteModelRequest
10
-from vllm.spec_decode.batch_expansion import BatchExpansionTop1Scorer
11
-from vllm.spec_decode.interfaces import SpeculativeProposals, SpeculativeScores
12
-from vllm.spec_decode.mqa_scorer import MQAScorer
13
-from vllm.worker.worker import Worker
14
-
15
-from .utils import create_batch, create_worker
16
-
17
-
18
-def create_proposal(propose_lens: list[int], vocab_size: int,
19
- device: str) -> SpeculativeProposals:
20
- batch_size = len(propose_lens)
21
- max_propose_len = max(propose_lens)
22
- proposal_probs = torch.rand((batch_size, max_propose_len, vocab_size),
23
- device=device)
24
-
25
- proposal_token_ids = torch.full((batch_size, max_propose_len),
26
- fill_value=-1,
27
- device=device)
28
- for i in range(batch_size):
29
- proposal_token_ids[i][:propose_lens[i]] = torch.argmax(
30
- proposal_probs[i][:propose_lens[i]], dim=-1)
31
-
32
- propose_lens = torch.tensor(propose_lens, device=device)
33
- return SpeculativeProposals(proposal_token_ids, proposal_probs,
34
- propose_lens)
35
-
36
-
37
-def assert_score_equal(score1: SpeculativeScores,
38
- score2: SpeculativeScores) -> None:
39
- assert torch.allclose(score1.probs, score2.probs)
40
- assert torch.allclose(score1.logprobs, score2.logprobs)
41
- assert torch.equal(
42
- score1.token_ids,
43
- score2.token_ids), f"{score1.token_ids}, {score2.token_ids}"
44
-
45
-
46
-@pytest.mark.parametrize('model_name', ['facebook/opt-125m'])
47
-@pytest.mark.parametrize('batch_size', [1, 2, 4, 8, 16])
48
-@pytest.mark.parametrize('max_propose_len', [1, 3, 5])
49
-@pytest.mark.parametrize('mixed_propose_len', [True])
50
-@pytest.mark.parametrize('device', ['cuda'])
51
-@pytest.mark.parametrize('prefill_chunking', [False, True])
52
-def test_scorer(model_name: str, batch_size: int, max_propose_len: int,
53
- mixed_propose_len: bool, device: str,
54
- prefill_chunking: bool) -> None:
55
- """
56
- Compare the batch expansion scorer and mqa scorer return the same score.
57
- We test for both queries with the same propose length and different
58
- propose length, as well as mixed prefill-decode batches.
59
- """
60
- seed = 0
61
- block_size = 32
62
- num_gpu_blocks = 2048 // block_size
63
- scorer_worker = create_worker(Worker, model_name, block_size,
64
- num_gpu_blocks, seed)
65
- scorer_worker.model_runner.disable_logprobs = True # accessed by mqa_scorer
66
- scorer_worker.model_runner.sampler.include_gpu_probs_tensor = True
67
- scorer_worker.model_runner.sampler.should_modify_greedy_probs_inplace = True
68
-
69
- vocab_size = scorer_worker.vocab_size
70
-
71
- if not mixed_propose_len:
72
- propose_lens = [max_propose_len] * batch_size
73
- else:
74
- # There must be at least 1 decode request, otherwise
75
- # we have nothing to score (`_run_no_spec`).
76
- non_zero_cnt = random.randint(1, batch_size)
77
- propose_lens = [max_propose_len
78
- ] * non_zero_cnt + [0] * (batch_size - non_zero_cnt)
79
- random.shuffle(propose_lens)
80
-
81
- seq_group_metadatalist, _, _ = create_batch(batch_size,
82
- max_propose_len,
83
- block_size=block_size,
84
- num_gpu_blocks=num_gpu_blocks)
85
-
86
- if mixed_propose_len and prefill_chunking and (n_prefills :=
87
- batch_size - non_zero_cnt):
88
- prefill, _, _ = create_batch(n_prefills,
89
- None,
90
- prefill_chunk_size=4,
91
- block_size=block_size,
92
- num_gpu_blocks=num_gpu_blocks,
93
- seq_ids=list(
94
- range(batch_size,
95
- batch_size + n_prefills)))
96
- # re-order to guarantee prefill|decode order
97
- target_group_metadatalist = [
98
- seq_group_metadatalist[i] for i, p in enumerate(propose_lens)
99
- if p > 0
100
- ]
101
- seq_group_metadatalist = prefill + target_group_metadatalist
102
- propose_lens = [0] * n_prefills + [p for p in propose_lens if p > 0]
103
-
104
- proposals = create_proposal(propose_lens, vocab_size, device)
105
- requests = ExecuteModelRequest(seq_group_metadatalist,
106
- num_lookahead_slots=max_propose_len)
107
-
108
- batch_expansion_scorer = BatchExpansionTop1Scorer(scorer_worker, device,
109
- vocab_size)
110
- batch_expansion_score = batch_expansion_scorer.score_proposals(
111
- requests, proposals)
112
-
113
- mqa_scorer = MQAScorer(scorer_worker, device, vocab_size)
114
- mqa_score = mqa_scorer.score_proposals(requests, proposals)
115
-
116
- assert_score_equal(batch_expansion_score, mqa_score)
tests/spec_decode/test_spec_decode_worker.py
deleted
-945
@@ -1,945 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import random
5
-from collections import defaultdict
6
-from types import SimpleNamespace
7
-from unittest.mock import MagicMock
8
-
9
-import pytest
10
-import torch
11
-
12
-from vllm.model_executor.layers.sampler import SamplerOutput
13
-from vllm.model_executor.utils import set_random_seed
14
-from vllm.sequence import ExecuteModelRequest, SequenceOutput
15
-from vllm.spec_decode.batch_expansion import BatchExpansionTop1Scorer
16
-from vllm.spec_decode.draft_model_runner import TP1DraftModelRunner
17
-from vllm.spec_decode.interfaces import SpeculativeProposals
18
-from vllm.spec_decode.metrics import (AsyncMetricsCollector,
19
- SpecDecodeWorkerMetrics)
20
-from vllm.spec_decode.multi_step_worker import MultiStepWorker
21
-from vllm.spec_decode.spec_decode_worker import (SpecDecodeWorker,
22
- split_num_cache_blocks_evenly)
23
-from vllm.worker.worker import Worker
24
-
25
-from .test_utils import mock_spec_decode_sampler
26
-from .utils import (create_batch, create_sampler_output_list, create_worker,
27
- mock_worker)
28
-
29
-
30
-@pytest.mark.parametrize('k', [1, 2, 6])
31
-@pytest.mark.parametrize('batch_size', [1, 2, 32])
32
-@pytest.mark.parametrize("acceptance_sampler_method",
33
- ["rejection_sampler", "typical_acceptance_sampler"])
34
-@torch.inference_mode()
35
-def test_correctly_calls_draft_model(k: int, batch_size: int,
36
- acceptance_sampler_method: str):
37
- """Verify SpecDecodeWorker calls the draft worker with correct
38
- inputs. Everything else is mocked out.
39
- """
40
- draft_worker = mock_worker(cls=MultiStepWorker)
41
- target_worker = mock_worker()
42
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
43
- worker = SpecDecodeWorker(
44
- draft_worker,
45
- target_worker,
46
- mock_spec_decode_sampler(acceptance_sampler_method),
47
- disable_logprobs=False,
48
- metrics_collector=metrics_collector)
49
- exception_secret = 'artificial stop'
50
- draft_worker.get_spec_proposals.side_effect = ValueError(exception_secret)
51
-
52
- seq_group_metadata_list, _, _ = create_batch(batch_size, k)
53
- execute_model_req = ExecuteModelRequest(
54
- seq_group_metadata_list=seq_group_metadata_list, num_lookahead_slots=k)
55
-
56
- with pytest.raises(ValueError, match=exception_secret):
57
- worker.execute_model(execute_model_req=execute_model_req)
58
-
59
- call_args_list = draft_worker.get_spec_proposals.call_args_list
60
- assert len(call_args_list) == 1
61
-
62
- for args, _ in call_args_list:
63
- actual_execute_model_data = args[0]
64
- assert actual_execute_model_data == execute_model_req
65
-
66
-
67
-@pytest.mark.parametrize('k', [1, 2, 6])
68
-@pytest.mark.parametrize('batch_size', [1, 2, 32])
69
-@pytest.mark.parametrize("acceptance_sampler_method",
70
- ["rejection_sampler", "typical_acceptance_sampler"])
71
-@torch.inference_mode()
72
-def test_batch_expansion_correctly_calls_target_model(
73
- k: int, batch_size: int, acceptance_sampler_method: str):
74
- """Verify SpecDecodeWorker calls the target model with correct
75
- inputs with batch expansion. Everything else is mocked out.
76
- """
77
- draft_worker = mock_worker(cls=MultiStepWorker, use_spec=False)
78
- target_worker = mock_worker(use_spec=False)
79
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
80
-
81
- draft_worker.device = 'cuda'
82
- target_worker.device = 'cuda'
83
-
84
- set_random_seed(1)
85
-
86
- worker = SpecDecodeWorker(
87
- draft_worker,
88
- target_worker,
89
- mock_spec_decode_sampler(acceptance_sampler_method),
90
- disable_logprobs=False,
91
- metrics_collector=metrics_collector,
92
- disable_mqa_scorer=True)
93
- worker.init_device()
94
-
95
- vocab_size = 32_000
96
-
97
- proposal_token_ids = torch.randint(low=0,
98
- high=vocab_size,
99
- size=(batch_size, k),
100
- dtype=torch.int64,
101
- device='cuda')
102
- proposal_probs = torch.rand(batch_size,
103
- k,
104
- vocab_size,
105
- dtype=torch.float32,
106
- device='cuda')
107
- proposal_lens = torch.ones(batch_size, dtype=torch.int64,
108
- device='cuda') * k
109
-
110
- seq_group_metadata_list, prompts, prev_output_tokens = create_batch(
111
- batch_size, k)
112
-
113
- draft_worker.get_spec_proposals.return_value = SpeculativeProposals(
114
- proposal_token_ids=proposal_token_ids,
115
- proposal_probs=proposal_probs,
116
- proposal_lens=proposal_lens)
117
-
118
- exception_secret = 'artificial stop'
119
- target_worker.execute_model.side_effect = ValueError(exception_secret)
120
-
121
- with pytest.raises(ValueError, match=exception_secret):
122
- worker.execute_model(execute_model_req=ExecuteModelRequest(
123
- seq_group_metadata_list=seq_group_metadata_list,
124
- num_lookahead_slots=k))
125
-
126
- seen_contexts: list[list[int]] = []
127
-
128
- call_args_list = target_worker.execute_model.call_args_list
129
- assert len(call_args_list) == 1
130
- for _, kwargs in call_args_list:
131
- seq_group_metadata_list = kwargs[
132
- "execute_model_req"].seq_group_metadata_list
133
-
134
- assert len(seq_group_metadata_list) == (k + 1) * batch_size
135
- for seq_group_metadata in seq_group_metadata_list:
136
- for seq_data in seq_group_metadata.seq_data.values():
137
- seen_contexts.append(seq_data.get_token_ids())
138
-
139
- expected_seen_contexts: list[list[int]] = []
140
-
141
- for prompt, prev_generated, draft_tokens in zip(
142
- prompts, prev_output_tokens, proposal_token_ids.tolist()):
143
-
144
- for i in range(len(draft_tokens) + 1):
145
- expected_seen_contexts.append(prompt + prev_generated +
146
- draft_tokens[:i])
147
-
148
- seen_contexts.sort()
149
- expected_seen_contexts.sort()
150
- assert expected_seen_contexts == seen_contexts
151
-
152
-
153
-@pytest.mark.parametrize('k', [1, 2, 6])
154
-@pytest.mark.parametrize('batch_size', [1, 2, 32])
155
-@pytest.mark.parametrize("acceptance_sampler_method",
156
- ["rejection_sampler", "typical_acceptance_sampler"])
157
-@torch.inference_mode()
158
-def test_correctly_calls_spec_decode_sampler(k: int, batch_size: int,
159
- acceptance_sampler_method: str):
160
- """Verify SpecDecodeWorker calls the rejection sampler with
161
- correct inputs. Everything else is mocked out.
162
- """
163
- vocab_size = 32_000
164
-
165
- draft_worker = mock_worker(cls=MultiStepWorker,
166
- vocab_size=vocab_size,
167
- use_spec=False)
168
- target_worker = mock_worker(vocab_size=vocab_size, use_spec=False)
169
- spec_decode_sampler = mock_spec_decode_sampler(acceptance_sampler_method)
170
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
171
- draft_worker.device = 'cuda'
172
- target_worker.device = 'cuda'
173
-
174
- set_random_seed(1)
175
-
176
- worker = SpecDecodeWorker(draft_worker,
177
- target_worker,
178
- spec_decode_sampler,
179
- disable_logprobs=False,
180
- metrics_collector=metrics_collector)
181
- worker.init_device()
182
-
183
- proposal_token_ids = torch.randint(low=0,
184
- high=vocab_size,
185
- size=(batch_size, k),
186
- dtype=torch.int64,
187
- device='cuda')
188
- proposal_probs = torch.rand(batch_size,
189
- k,
190
- vocab_size,
191
- dtype=torch.float32,
192
- device='cuda')
193
-
194
- proposal_lens = torch.ones(batch_size, dtype=torch.int64,
195
- device='cuda') * k
196
-
197
- seq_group_metadata_list, _, _ = create_batch(batch_size, k)
198
-
199
- draft_worker.get_spec_proposals.return_value = SpeculativeProposals(
200
- proposal_token_ids=proposal_token_ids,
201
- proposal_probs=proposal_probs,
202
- proposal_lens=proposal_lens)
203
-
204
- target_token_ids = torch.randint(low=0,
205
- high=vocab_size,
206
- size=(1, batch_size * (k + 1)),
207
- dtype=torch.int64,
208
- device='cuda')
209
- target_token_probs = torch.rand(1,
210
- batch_size * (k + 1),
211
- vocab_size,
212
- dtype=torch.float32,
213
- device='cuda')
214
- target_token_logprobs = torch.rand(1,
215
- batch_size * (k + 1),
216
- vocab_size,
217
- dtype=torch.float32,
218
- device='cuda')
219
- target_output = create_sampler_output_list(target_token_ids,
220
- target_token_probs,
221
- target_token_logprobs)
222
-
223
- target_worker.execute_model.return_value = [target_output[0]]
224
-
225
- exception_secret = 'artificial stop'
226
-
227
- spec_decode_sampler.side_effect = ValueError(exception_secret)
228
-
229
- with pytest.raises(ValueError, match=exception_secret):
230
- worker.execute_model(execute_model_req=ExecuteModelRequest(
231
- seq_group_metadata_list=seq_group_metadata_list,
232
- num_lookahead_slots=k))
233
-
234
- assert len(spec_decode_sampler.call_args_list) == 1
235
- _, kwargs = spec_decode_sampler.call_args_list[0]
236
- actual = SimpleNamespace(**kwargs)
237
-
238
- assert torch.equal(actual.bonus_token_ids,
239
- target_token_ids.reshape(batch_size, k + 1)[:, -1:])
240
- assert torch.equal(actual.target_with_bonus_probs,
241
- target_token_probs.reshape(batch_size, k + 1, -1))
242
- assert torch.equal(actual.draft_token_ids, proposal_token_ids)
243
- assert torch.equal(actual.draft_probs, proposal_probs)
244
-
245
-
246
-@pytest.mark.parametrize('k', [1, 2, 6])
247
-@pytest.mark.parametrize('batch_size', [1, 2, 32])
248
-@pytest.mark.parametrize("acceptance_sampler_method",
249
- ["rejection_sampler", "typical_acceptance_sampler"])
250
-@torch.inference_mode()
251
-def test_correctly_formats_output(k: int, batch_size: int,
252
- acceptance_sampler_method: str):
253
- """Verify SpecDecodeWorker formats sampler output correctly.
254
- Everything else is mocked out.
255
- """
256
- vocab_size = 32_000
257
-
258
- draft_worker = mock_worker(cls=MultiStepWorker,
259
- vocab_size=vocab_size,
260
- use_spec=False)
261
- target_worker = mock_worker(vocab_size=vocab_size, use_spec=False)
262
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
263
- draft_worker.device = 'cuda'
264
- target_worker.device = 'cuda'
265
-
266
- set_random_seed(1)
267
- spec_decode_sampler = mock_spec_decode_sampler(acceptance_sampler_method)
268
- worker = SpecDecodeWorker(draft_worker,
269
- target_worker,
270
- spec_decode_sampler,
271
- disable_logprobs=False,
272
- metrics_collector=metrics_collector)
273
- worker.init_device()
274
-
275
- proposal_token_ids = torch.randint(low=0,
276
- high=vocab_size,
277
- size=(batch_size, k),
278
- dtype=torch.int64,
279
- device='cuda')
280
- proposal_probs = torch.rand(batch_size,
281
- k,
282
- vocab_size,
283
- dtype=torch.float32,
284
- device='cuda')
285
-
286
- proposal_lens = torch.ones(batch_size, dtype=torch.int64,
287
- device='cuda') * k
288
-
289
- seq_group_metadata_list, _, _ = create_batch(batch_size, k)
290
-
291
- draft_worker.get_spec_proposals.return_value = SpeculativeProposals(
292
- proposal_token_ids=proposal_token_ids,
293
- proposal_probs=proposal_probs,
294
- proposal_lens=proposal_lens)
295
-
296
- target_token_ids = torch.randint(low=0,
297
- high=vocab_size,
298
- size=(1, batch_size * (k + 1)),
299
- dtype=torch.int64,
300
- device='cuda')
301
- target_token_probs = torch.rand(1,
302
- batch_size * (k + 1),
303
- vocab_size,
304
- dtype=torch.float32,
305
- device='cuda')
306
- target_token_logprobs = torch.rand(1,
307
- batch_size * (k + 1),
308
- vocab_size,
309
- dtype=torch.float32,
310
- device='cuda')
311
- target_output = create_sampler_output_list(target_token_ids,
312
- target_token_probs,
313
- target_token_logprobs)
314
-
315
- target_worker.execute_model.return_value = [target_output[0]]
316
-
317
- spec_decode_sampler_output = torch.randint(low=0,
318
- high=vocab_size,
319
- size=(batch_size, k + 1),
320
- dtype=torch.int64,
321
- device='cuda')
322
- for i in range(batch_size):
323
- minimum_accepted_tokens = 1
324
- spec_decode_sampler_output[i][
325
- -random.randint(minimum_accepted_tokens, k + 1):] = -1
326
-
327
- spec_decode_sampler.return_value = spec_decode_sampler_output
328
- output = worker.execute_model(execute_model_req=ExecuteModelRequest(
329
- seq_group_metadata_list=seq_group_metadata_list,
330
- num_lookahead_slots=k))
331
-
332
- expected_output = create_sampler_output_list(
333
- token_ids=spec_decode_sampler_output.transpose(0, 1),
334
- probs=[None for _ in range(k + 1)],
335
- logprobs=[None for _ in range(k + 1)])
336
-
337
- seq_ids = [
338
- next(iter(seq_group_metadata.seq_data.keys()))
339
- for seq_group_metadata in seq_group_metadata_list
340
- ]
341
- actual_output_by_seq: dict[int, list[SequenceOutput]] = {
342
- seq_id: []
343
- for seq_id in seq_ids
344
- }
345
- expected_output_by_seq: dict[int, list[SequenceOutput]] = {
346
- seq_id: []
347
- for seq_id in seq_ids
348
- }
349
-
350
- for step in output:
351
- for seq_group in step:
352
- for sample in seq_group.samples:
353
- seq_id = sample.parent_seq_id
354
- actual_output_by_seq[seq_id].append(sample)
355
-
356
- for step in expected_output:
357
- for seq_group in step:
358
- for sample in seq_group.samples:
359
- seq_id = sample.parent_seq_id
360
- expected_output_by_seq[seq_id].append(sample)
361
-
362
- all_seen_seq_ids = set(
363
- list(actual_output_by_seq.keys()) +
364
- list(expected_output_by_seq.keys()))
365
- for seq_id in all_seen_seq_ids:
366
- actual_by_step = actual_output_by_seq[seq_id]
367
- expected_by_step = expected_output_by_seq[seq_id]
368
-
369
- for i in range(k + 1):
370
- if i >= len(actual_by_step):
371
- assert expected_by_step[i].output_token == -1
372
- continue
373
- assert actual_by_step[i].output_token == expected_by_step[
374
- i].output_token
375
-
376
-
377
-@pytest.mark.parametrize('k', [1, 2])
378
-@pytest.mark.parametrize('batch_size', [1])
379
-@pytest.mark.parametrize('returns_metrics', [True, False])
380
-@pytest.mark.parametrize("acceptance_sampler_method",
381
- ["rejection_sampler", "typical_acceptance_sampler"])
382
-@torch.inference_mode()
383
-def test_collects_metrics(k: int, batch_size: int, returns_metrics: bool,
384
- acceptance_sampler_method: str):
385
- """Verify SpecDecodeWorker collects metrics.
386
- """
387
- vocab_size = 32_000
388
-
389
- draft_worker = mock_worker(cls=MultiStepWorker,
390
- vocab_size=vocab_size,
391
- use_spec=False)
392
- target_worker = mock_worker(vocab_size=vocab_size, use_spec=False)
393
- spec_decode_sampler = mock_spec_decode_sampler(acceptance_sampler_method)
394
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
395
- draft_worker.device = 'cuda'
396
- target_worker.device = 'cuda'
397
-
398
- set_random_seed(1)
399
-
400
- worker = SpecDecodeWorker(draft_worker,
401
- target_worker,
402
- spec_decode_sampler,
403
- disable_logprobs=False,
404
- metrics_collector=metrics_collector)
405
- worker.init_device()
406
-
407
- proposal_token_ids = torch.randint(low=0,
408
- high=vocab_size,
409
- size=(batch_size, k),
410
- dtype=torch.int64,
411
- device='cuda')
412
- proposal_probs = torch.rand(batch_size,
413
- k,
414
- vocab_size,
415
- dtype=torch.float32,
416
- device='cuda')
417
-
418
- proposal_lens = torch.ones(batch_size, dtype=torch.int64,
419
- device='cuda') * k
420
-
421
- seq_group_metadata_list, _, _ = create_batch(batch_size, k)
422
-
423
- draft_worker.get_spec_proposals.return_value = SpeculativeProposals(
424
- proposal_token_ids=proposal_token_ids,
425
- proposal_probs=proposal_probs,
426
- proposal_lens=proposal_lens)
427
-
428
- target_token_ids = torch.randint(low=0,
429
- high=vocab_size,
430
- size=(1, batch_size * (k + 1)),
431
- dtype=torch.int64,
432
- device='cuda')
433
- target_token_probs = torch.rand(1,
434
- batch_size * (k + 1),
435
- vocab_size,
436
- dtype=torch.float32,
437
- device='cuda')
438
- target_token_logprobs = torch.rand(1,
439
- batch_size * (k + 1),
440
- vocab_size,
441
- dtype=torch.float32,
442
- device='cuda')
443
- target_output = create_sampler_output_list(target_token_ids,
444
- target_token_probs,
445
- target_token_logprobs)
446
-
447
- target_worker.execute_model.return_value = [target_output[0]]
448
-
449
- spec_decode_sampler_output = torch.randint(low=0,
450
- high=vocab_size,
451
- size=(batch_size, k + 1),
452
- dtype=torch.int64,
453
- device='cuda')
454
- for i in range(batch_size):
455
- minimum_accepted_tokens = 1
456
- spec_decode_sampler_output[i][
457
- -random.randint(minimum_accepted_tokens, k + 1):] = -1
458
- spec_decode_sampler.return_value = spec_decode_sampler_output
459
-
460
- mock_rejsample_metrics = MagicMock(
461
- spec=SpecDecodeWorkerMetrics) if returns_metrics else None
462
- metrics_collector.maybe_collect_rejsample_metrics.return_value = (
463
- mock_rejsample_metrics)
464
-
465
- output = worker.execute_model(execute_model_req=ExecuteModelRequest(
466
- seq_group_metadata_list=seq_group_metadata_list,
467
- num_lookahead_slots=k))
468
- assert output[0].spec_decode_worker_metrics == mock_rejsample_metrics
469
-
470
- call_args_list = (
471
- metrics_collector.maybe_collect_rejsample_metrics.call_args_list)
472
- assert len(call_args_list) == 1
473
- args, kwargs = call_args_list[0]
474
- assert args[0] == k or kwargs.get('k', -1) == k
475
-
476
-
477
-@pytest.mark.parametrize('k', [0])
478
-@pytest.mark.parametrize('batch_size', [1, 2, 32])
479
-@pytest.mark.parametrize("acceptance_sampler_method",
480
- ["rejection_sampler", "typical_acceptance_sampler"])
481
-@torch.inference_mode()
482
-def test_k_equals_zero(k: int, batch_size: int,
483
- acceptance_sampler_method: str):
484
- """Verify that the SpecDecodeWorker calls the draft and target workers
485
- when k is zero. This happens during prefill.
486
- """
487
- draft_worker = mock_worker(cls=MultiStepWorker)
488
- target_worker = mock_worker()
489
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
490
-
491
- sampler_output = MagicMock(spec=SamplerOutput)
492
- sampler_output.hidden_states = None
493
- target_worker.execute_model.return_value = [sampler_output]
494
-
495
- draft_worker.device = 'cuda'
496
- target_worker.device = 'cuda'
497
-
498
- set_random_seed(1)
499
-
500
- worker = SpecDecodeWorker(
501
- proposer_worker=draft_worker,
502
- scorer_worker=target_worker,
503
- spec_decode_sampler=mock_spec_decode_sampler(
504
- acceptance_sampler_method),
505
- disable_logprobs=False,
506
- metrics_collector=metrics_collector,
507
- )
508
-
509
- seq_group_metadata_list, _, _ = create_batch(batch_size,
510
- k,
511
- prev_output_token_len=0)
512
- execute_model_req = ExecuteModelRequest(
513
- seq_group_metadata_list=seq_group_metadata_list, num_lookahead_slots=k)
514
-
515
- out = worker.execute_model(execute_model_req=execute_model_req)
516
-
517
- assert len(out) == 1, f"expected only one token output when {k=}"
518
- assert out[0].sampled_token_probs is None, (
519
- "expect gpu tensor references to be None")
520
- assert out[
521
- 0].sampled_token_ids is None, "expect gpu tensor references to be None"
522
-
523
- draft_worker.execute_model.assert_called_once_with(execute_model_req)
524
- target_worker.execute_model.assert_called_once_with(execute_model_req)
525
-
526
-
527
-@pytest.mark.parametrize('k', [0, 5])
528
-@pytest.mark.parametrize('batch_size', [0])
529
-@pytest.mark.parametrize("acceptance_sampler_method",
530
- ["rejection_sampler", "typical_acceptance_sampler"])
531
-@torch.inference_mode()
532
-def test_empty_input_batch(k: int, batch_size: int,
533
- acceptance_sampler_method: str):
534
- """Verify that the SpecDecodeWorker calls the draft and target workers
535
- when the input batch is empty. This can happen if the engine communicates
536
- to the workers information without scheduling a batch.
537
- """
538
- draft_worker = mock_worker(cls=MultiStepWorker)
539
- target_worker = mock_worker()
540
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
541
-
542
- sampler_output = MagicMock(spec=SamplerOutput)
543
- sampler_output.hidden_states = None
544
- target_worker.execute_model.return_value = [sampler_output]
545
-
546
- draft_worker.device = 'cuda'
547
- target_worker.device = 'cuda'
548
-
549
- set_random_seed(1)
550
-
551
- worker = SpecDecodeWorker(
552
- proposer_worker=draft_worker,
553
- scorer_worker=target_worker,
554
- spec_decode_sampler=mock_spec_decode_sampler(
555
- acceptance_sampler_method),
556
- disable_logprobs=False,
557
- metrics_collector=metrics_collector,
558
- )
559
-
560
- seq_group_metadata_list, _, _ = create_batch(batch_size,
561
- k,
562
- prev_output_token_len=0)
563
- execute_model_req = ExecuteModelRequest(
564
- seq_group_metadata_list=seq_group_metadata_list, num_lookahead_slots=k)
565
-
566
- out = worker.execute_model(execute_model_req=execute_model_req)
567
-
568
- assert len(out) == 1, f"expected only one token output when {k=}"
569
- assert out[0].sampled_token_probs is None, (
570
- "expect gpu tensor references to be None")
571
- assert out[
572
- 0].sampled_token_ids is None, "expect gpu tensor references to be None"
573
-
574
- draft_worker.execute_model.assert_called_once_with(execute_model_req)
575
- target_worker.execute_model.assert_called_once_with(execute_model_req)
576
-
577
-
578
-@pytest.mark.parametrize("acceptance_sampler_method",
579
- ["rejection_sampler", "typical_acceptance_sampler"])
580
-@pytest.mark.skip_global_cleanup
581
-def test_init_device(acceptance_sampler_method: str):
582
- """Verify SpecDecodeWorker invokes proposer/scorer worker init_device, as
583
- well as other GPU initialization.
584
- """
585
- draft_worker = mock_worker(cls=MultiStepWorker, use_spec=False)
586
- target_worker = mock_worker(use_spec=False)
587
- spec_decode_sampler = mock_spec_decode_sampler(acceptance_sampler_method)
588
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
589
-
590
- worker = SpecDecodeWorker(
591
- proposer_worker=draft_worker,
592
- scorer_worker=target_worker,
593
- spec_decode_sampler=spec_decode_sampler,
594
- disable_logprobs=False,
595
- metrics_collector=metrics_collector,
596
- )
597
- worker.init_device()
598
-
599
- draft_worker.init_device.assert_called_once()
600
-
601
- target_worker.init_device.assert_called_once()
602
-
603
- metrics_collector.init_tensors.assert_called_once()
604
- spec_decode_sampler.init_tensors.assert_called_once()
605
-
606
-
607
-@pytest.mark.parametrize("acceptance_sampler_method",
608
- ["rejection_sampler", "typical_acceptance_sampler"])
609
-@torch.inference_mode()
610
-def test_initialize_cache(acceptance_sampler_method):
611
- """Verify SpecDecodeWorker invokes initialize_cache on proposer/scorer
612
- workers.
613
- """
614
- draft_worker = mock_worker(cls=MultiStepWorker)
615
- target_worker = mock_worker()
616
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
617
-
618
- worker = SpecDecodeWorker(proposer_worker=draft_worker,
619
- scorer_worker=target_worker,
620
- spec_decode_sampler=mock_spec_decode_sampler(
621
- acceptance_sampler_method),
622
- metrics_collector=metrics_collector)
623
-
624
- kwargs = {"num_gpu_blocks": 1024, "num_cpu_blocks": 1023}
625
- worker.initialize_cache(**kwargs)
626
-
627
- draft_worker.initialize_cache.assert_called_once_with(**kwargs)
628
- target_worker.initialize_cache.assert_called_once_with(**kwargs)
629
-
630
-
631
-@pytest.mark.parametrize('available_gpu_blocks', [1, 1024])
632
-@pytest.mark.parametrize('available_cpu_blocks', [500])
633
-@pytest.mark.parametrize('target_cache_block_size_bytes', [2 * 2 * 4096])
634
-@pytest.mark.parametrize('draft_kv_size_bytes', [0, 2 * 2 * 768, 2 * 2 * 4096])
635
-@pytest.mark.parametrize("acceptance_sampler_method",
636
- ["rejection_sampler", "typical_acceptance_sampler"])
637
-@pytest.mark.skip_global_cleanup
638
-def test_determine_num_available_blocks(available_gpu_blocks: int,
639
- available_cpu_blocks: int,
640
- target_cache_block_size_bytes: int,
641
- draft_kv_size_bytes: int,
642
- acceptance_sampler_method: str):
643
- """Verify SpecDecodeWorker correctly profiles num available GPU blocks.
644
- Specifically, it should run profiling in the scorer worker, and then evenly
645
- split the blocks between proposer and scorer worker.
646
- """
647
- draft_worker = mock_worker(cls=MultiStepWorker)
648
- target_worker = mock_worker()
649
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
650
-
651
- target_worker.determine_num_available_blocks.return_value = (
652
- available_gpu_blocks, available_cpu_blocks)
653
- target_worker.get_cache_block_size_bytes.return_value = (
654
- target_cache_block_size_bytes)
655
- draft_worker.get_cache_block_size_bytes.return_value = draft_kv_size_bytes
656
-
657
- worker = SpecDecodeWorker(
658
- draft_worker, target_worker,
659
- mock_spec_decode_sampler(acceptance_sampler_method), metrics_collector)
660
-
661
- num_gpu_blocks, num_cpu_blocks = worker.determine_num_available_blocks()
662
-
663
- target_worker.determine_num_available_blocks.assert_called_once()
664
- assert num_cpu_blocks == available_cpu_blocks
665
-
666
- assert num_gpu_blocks == split_num_cache_blocks_evenly(
667
- target_cache_block_size_bytes, draft_kv_size_bytes,
668
- available_gpu_blocks)
669
-
670
-
671
-@pytest.mark.parametrize('available_gpu_blocks',
672
- list(range(20)) + [1024, 1024**2])
673
-@pytest.mark.parametrize('target_cache_block_size_bytes',
674
- [2 * 2 * 4096, 2 * 2 * 8192])
675
-@pytest.mark.parametrize('draft_kv_size_bytes', [0, 2 * 2 * 768, 2 * 2 * 4096])
676
-@pytest.mark.skip_global_cleanup
677
-def test_split_num_cache_blocks_evenly(available_gpu_blocks: int,
678
- target_cache_block_size_bytes: int,
679
- draft_kv_size_bytes: int):
680
- """Verify split_num_cache_blocks_evenly does not exceed original memory
681
- allocation in bytes.
682
- """
683
- num_blocks = split_num_cache_blocks_evenly(target_cache_block_size_bytes,
684
- draft_kv_size_bytes,
685
- available_gpu_blocks)
686
- assert (num_blocks * target_cache_block_size_bytes) + (
687
- num_blocks * draft_kv_size_bytes) <= (available_gpu_blocks *
688
- target_cache_block_size_bytes)
689
-
690
-
691
-@torch.inference_mode()
692
-def test_populate_seq_ids_with_bonus_tokens():
693
- """
694
- Verify that a call to _create_output_sampler_list correctly updates
695
- seq_with_bonus_token_in_last_step.
696
-
697
- seq_with_bonus_token_in_last_step is an internal data structure in
698
- SpecDecodeWorker that tracks the sequence IDs which are assigned bonus
699
- tokens by the target model in their last forward pass. This state is
700
- maintained only for models relying on the KV cache, such as those using
701
- the MultiStepWorker.
702
- """
703
- batch_size = 10
704
- k = 5
705
- vocab_size = 10000
706
- num_sequences_with_bonus_tokens = 5
707
- target_worker = mock_worker(vocab_size=vocab_size, use_spec=False)
708
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
709
- target_worker.execute_model.return_value = [MagicMock(spec=SamplerOutput)]
710
- target_worker.device = 'cuda'
711
-
712
- set_random_seed(1)
713
- draft_worker = mock_worker(cls=MultiStepWorker)
714
- draft_worker.device = 'cuda'
715
- # The sequence_ids attached to each sequence in the batch.
716
- # The sequence at index i has seq_id assigned_seq_ids[i]
717
- assigned_seq_ids = list(range(batch_size))
718
- seq_group_metadata_list, _, _ = create_batch(batch_size,
719
- k,
720
- seq_ids=assigned_seq_ids,
721
- prev_output_token_len=10)
722
- target_token_logprobs = torch.rand(batch_size, (k + 1),
723
- vocab_size,
724
- dtype=torch.float32,
725
- device='cuda')
726
- accepted_token_ids = torch.randint(low=0,
727
- high=vocab_size,
728
- size=(batch_size, (k + 1)),
729
- dtype=torch.int64,
730
- device='cuda')
731
- expected_request_id_seq_ids_mapping: dict[str, set[int]] = defaultdict(set)
732
- for seq_group_metadata in seq_group_metadata_list:
733
- for seq_id in seq_group_metadata.seq_data:
734
- expected_request_id_seq_ids_mapping[
735
- seq_group_metadata.request_id].add(seq_id)
736
- # Generate a random sample of sequence indexes with bonus tokens
737
- seq_indexes_with_bonus_tokens = random.sample(
738
- range(batch_size), num_sequences_with_bonus_tokens)
739
- # Create a mask that is True for indices in seq_indexes_with_bonus_tokens
740
- mask = torch.ones(batch_size, dtype=torch.bool, device='cuda')
741
- mask[seq_indexes_with_bonus_tokens] = False
742
- # Set the last token ID to -1 for all indices not in
743
- # seq_indexes_with_bonus_tokens to indicate the lack of bonus token in
744
- # those indices.
745
- accepted_token_ids[mask, -1:] = -1
746
- worker = SpecDecodeWorker(draft_worker,
747
- target_worker,
748
- mock_spec_decode_sampler("rejection_sampler"),
749
- disable_logprobs=False,
750
- metrics_collector=metrics_collector)
751
- # Initialize _seq_with_bonus_token_in_last_step with a set of sequence IDs.
752
- # This set includes all sequence IDs in the batch as well as an additional
753
- # `num_extra_sequence_ids` sequence IDs. Note that the sequence IDs are in
754
- # the range [0, batch_size + num_extra_sequence_ids).
755
- num_extra_sequence_ids = 10
756
- worker._seq_with_bonus_token_in_last_step = set(
757
- range(batch_size + num_extra_sequence_ids))
758
- worker._create_output_sampler_list(
759
- seq_group_metadata_list=seq_group_metadata_list,
760
- accepted_token_ids=accepted_token_ids,
761
- target_logprobs=target_token_logprobs,
762
- prompt_logprobs=None,
763
- k=k,
764
- stage_times=(0, 0, 0))
765
- # Verify that _seq_with_bonus_token_in_last_step contains the following:
766
- # 1. Sequence IDs that were already present in
767
- # _seq_with_bonus_token_in_last_step but were not part of the current
768
- # batch are retained.
769
- # 2. Of the sequence IDs present in the current batch, only those with a
770
- # bonus token are retained in _seq_with_bonus_token_in_last_step.
771
- # Sequence IDs that are present in the current batch but do not have
772
- # bonus tokens are removed from _seq_with_bonus_token_in_last_step.
773
- expected_seq_ids_with_bonus_tokens = \
774
- set([assigned_seq_ids[i] for i in seq_indexes_with_bonus_tokens])
775
- additional_sequence_ids = \
776
- set(range(batch_size, batch_size + num_extra_sequence_ids))
777
- assert worker._seq_with_bonus_token_in_last_step == \
778
- expected_seq_ids_with_bonus_tokens.union(additional_sequence_ids)
779
- assert worker._request_id_seq_id_mapping == \
780
- expected_request_id_seq_ids_mapping
781
-
782
-
783
-@torch.inference_mode()
784
-def test_handle_finished_requests():
785
- """
786
- Test to verify that finished request IDs are appropriately processed to
787
- update the internal state of the SpecDecodeWorker.
788
-
789
- This test initializes the SpecDecodeWorker with mock data, marks certain
790
- requests as finished, and ensures that the corresponding sequence IDs are
791
- correctly removed from the internal mappings.
792
- """
793
- batch_size = 32
794
- k = 3
795
- draft_worker = mock_worker(cls=MultiStepWorker)
796
- target_worker = mock_worker()
797
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
798
- worker = SpecDecodeWorker(draft_worker, target_worker,
799
- mock_spec_decode_sampler("rejection_sampler"),
800
- metrics_collector)
801
- # Initialize the request_id_seq_id_mapping mapping dict with a few fake
802
- # request ids and corresponding sequence ids.
803
- worker._request_id_seq_id_mapping = \
804
- {'request-1': {1,2,3}, 'request-2': {4,5,6,7},
805
- 'request-3': {8,9}, 'request-4': {10,11}}
806
- # Initialize seq_with_bonus_token_in_last_step with a few fake
807
- # sequence ids.
808
- worker._seq_with_bonus_token_in_last_step = {1, 4, 5, 8, 9, 10}
809
- exception_secret = 'artificial stop'
810
- draft_worker.get_spec_proposals.side_effect = ValueError(exception_secret)
811
-
812
- seq_group_metadata_list, _, _ = create_batch(batch_size, k)
813
- # Mark requests with ids request-1 and request-3 as finished.
814
- execute_model_req = ExecuteModelRequest(
815
- seq_group_metadata_list=seq_group_metadata_list,
816
- num_lookahead_slots=k,
817
- finished_requests_ids=['request-1', 'request-3'])
818
-
819
- with pytest.raises(ValueError, match=exception_secret):
820
- worker.execute_model(execute_model_req=execute_model_req)
821
- # Verify that request-1 and request-3 are removed from
822
- # request_id_seq_id_mapping
823
- assert worker._request_id_seq_id_mapping == \
824
- {'request-2': {4,5,6,7}, 'request-4': {10,11}}
825
- # Verify that all sequence ids corresponding to 'request-1'
826
- # and 'request-3' are removed from seq_with_bonus_token_in_last_step.
827
- assert worker._seq_with_bonus_token_in_last_step == \
828
- {4,5,10}
829
-
830
-
831
-@pytest.mark.parametrize('k', [3])
832
-@pytest.mark.parametrize('batch_size', [2, 32])
833
-@pytest.mark.parametrize("batch_composition",
834
- ["prefill_only", "decode_only", "mixed"])
835
-@torch.inference_mode()
836
-def test_chunked_prefill_flow(k: int, batch_size: int, batch_composition: str):
837
- """
838
- Verify SpecDecodeWorker calls match the expected flow.
839
- """
840
- vocab_size = 32_000
841
- draft_worker = mock_worker(cls=MultiStepWorker)
842
- target_worker = mock_worker()
843
- metrics_collector = MagicMock(spec=AsyncMetricsCollector)
844
- worker = SpecDecodeWorker(draft_worker,
845
- target_worker,
846
- mock_spec_decode_sampler("rejection_sampler"),
847
- disable_logprobs=False,
848
- metrics_collector=metrics_collector)
849
- exception_secret = 'artificial stop'
850
- worker.scorer = mock_worker(BatchExpansionTop1Scorer)
851
- worker.scorer.score_proposals.side_effect = ValueError(exception_secret)
852
-
853
- # Create batch with combination of terminal/non-terminal prefill chunks
854
- # and decodes (different seq_ids).
855
- decodes, _, _ = create_batch(batch_size, k)
856
- # Pre-chunking here, get 'batch_size' chunks.
857
- prefill, _, _ = create_batch(batch_size,
858
- k,
859
- prefill_chunk_size=4,
860
- seq_ids=list(range(batch_size,
861
- batch_size * 2)))
862
-
863
- if batch_composition == "prefill_only":
864
- n_prefills = batch_size
865
- elif batch_composition == "decode_only":
866
- n_prefills = 0
867
- else:
868
- n_prefills = random.randint(1, batch_size - 1)
869
- n_decodes = batch_size - n_prefills
870
-
871
- prefill = random.sample(prefill, n_prefills)
872
- decodes = random.sample(decodes, n_decodes)
873
- target_group_metadata_list = prefill + decodes
874
- execute_model_req = ExecuteModelRequest(
875
- seq_group_metadata_list=target_group_metadata_list,
876
- # For prefill only batches we expect num_lookahead_slots = 0.
877
- num_lookahead_slots=k if n_decodes > 0 else 0)
878
-
879
- target_token_ids = torch.randint(low=0,
880
- high=vocab_size,
881
- size=(1, batch_size * (k + 1)),
882
- dtype=torch.int64,
883
- device='cuda')
884
- target_token_probs = torch.rand(1,
885
- batch_size * (k + 1),
886
- vocab_size,
887
- dtype=torch.float32,
888
- device='cuda')
889
- target_token_logprobs = torch.rand(1,
890
- batch_size * (k + 1),
891
- vocab_size,
892
- dtype=torch.float32,
893
- device='cuda')
894
- target_output = create_sampler_output_list(target_token_ids,
895
- target_token_probs,
896
- target_token_logprobs)
897
-
898
- target_worker.execute_model.return_value = [target_output[0]]
899
-
900
- if not len(decodes):
901
- worker.execute_model(execute_model_req=execute_model_req)
902
- # no spec run (prefill only)
903
- draft_worker.execute_model.assert_called_once_with(execute_model_req)
904
- target_worker.execute_model.assert_called_once_with(execute_model_req)
905
- else:
906
- # Decode-only run OR mixed batch, scorer call fails (it's mocked)
907
- with pytest.raises(ValueError, match=exception_secret):
908
- worker.execute_model(execute_model_req=execute_model_req)
909
- # but first draft still counted
910
- assert draft_worker.get_spec_proposals.call_count == 1
911
-
912
-
913
-def test_correctly_load_weight_for_eagle():
914
- """
915
- Verify SpecDecodeWorker loads lm_head weight for eagle correctly.
916
- """
917
- seed = 100
918
- block_size = 32
919
- num_gpu_blocks = 8096 // block_size
920
- target_worker = create_worker(
921
- Worker,
922
- "JackFram/llama-68m",
923
- block_size,
924
- num_gpu_blocks,
925
- seed,
926
- )
927
- draft_worker = create_worker(
928
- MultiStepWorker,
929
- "abhigoyal/vllm-eagle-llama-68m-random",
930
- block_size,
931
- num_gpu_blocks,
932
- seed,
933
- model_runner_cls=TP1DraftModelRunner,
934
- )
935
-
936
- spec_decode_sampler = mock_spec_decode_sampler("rejection_sampler")
937
- worker = SpecDecodeWorker(draft_worker,
938
- target_worker,
939
- spec_decode_sampler,
940
- disable_logprobs=False)
941
- worker.proposer_worker.maybe_load_lm_head_weight(
942
- target_worker.model_runner.model.lm_head.weight.data)
943
- assert torch.allclose(
944
- worker.proposer_worker.worker.model_runner.model.lm_head.weight.data,
945
- worker.scorer_worker.model_runner.model.lm_head.weight.data)
tests/spec_decode/test_utils.py
deleted
-150
@@ -1,150 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from unittest.mock import MagicMock
5
-
6
-import pytest
7
-import torch
8
-
9
-from vllm.model_executor.layers.rejection_sampler import RejectionSampler
10
-from vllm.model_executor.layers.sampler import _get_ranks
11
-from vllm.model_executor.layers.typical_acceptance_sampler import (
12
- TypicalAcceptanceSampler)
13
-from vllm.sequence import SequenceGroupMetadata, get_all_seq_ids
14
-from vllm.spec_decode.util import (get_sampled_token_logprobs,
15
- split_batch_by_proposal_len)
16
-
17
-
18
-def test_get_all_seq_ids():
19
- """Verify get_all_seq_ids extracts all seq ids.
20
- """
21
- expected_seq_ids = list(range(10)) + list(range(100, 110))
22
-
23
- seq_group_metadata_list = [
24
- SequenceGroupMetadata(
25
- request_id=str(seq_id),
26
- is_prompt=True,
27
- seq_data={
28
- seq_id: MagicMock(),
29
- },
30
- sampling_params=MagicMock(),
31
- block_tables={
32
- seq_id: MagicMock(),
33
- },
34
- lora_request=None,
35
- ) for seq_id in expected_seq_ids
36
- ]
37
-
38
- actual_seq_ids = get_all_seq_ids(seq_group_metadata_list)
39
- assert actual_seq_ids == expected_seq_ids
40
-
41
-
42
-@pytest.fixture
43
-def fake_sequence_group_metadata():
44
- seq_ids = list(range(3))
45
- return [
46
- SequenceGroupMetadata(
47
- request_id=str(i),
48
- is_prompt=True,
49
- seq_data={
50
- i: MagicMock(),
51
- },
52
- sampling_params=MagicMock(),
53
- block_tables={
54
- i: MagicMock(),
55
- },
56
- lora_request=None,
57
- ) for i in seq_ids
58
- ]
59
-
60
-
61
-def test_filter_zero_length_proposals(fake_sequence_group_metadata):
62
- proposal_lens = [0, 1, 0]
63
- _, (filtered_groups,
64
- indices) = split_batch_by_proposal_len(fake_sequence_group_metadata,
65
- proposal_lens)
66
-
67
- expected_groups = [
68
- fake_sequence_group_metadata[0], fake_sequence_group_metadata[2]
69
- ]
70
- expected_indices = [0, 2]
71
-
72
- assert filtered_groups == expected_groups
73
- assert indices == expected_indices
74
-
75
-
76
-def test_filter_non_zero_length_proposals(fake_sequence_group_metadata):
77
- proposal_lens = [0, 1, 2]
78
- (filtered_groups,
79
- indices), _ = split_batch_by_proposal_len(fake_sequence_group_metadata,
80
- proposal_lens)
81
-
82
- expected_groups = [
83
- fake_sequence_group_metadata[1], fake_sequence_group_metadata[2]
84
- ]
85
- expected_indices = [1, 2]
86
-
87
- assert filtered_groups == expected_groups
88
- assert indices == expected_indices
89
-
90
-
91
-def test_empty_inputs():
92
- _, (filtered_groups, indices) = split_batch_by_proposal_len([], [])
93
-
94
- assert filtered_groups == []
95
- assert indices == []
96
-
97
-
98
-def test_all_zero_with_non_zero_filter(fake_sequence_group_metadata):
99
- proposal_lens = [0, 0, 0]
100
- (filtered_groups,
101
- indices), _ = split_batch_by_proposal_len(fake_sequence_group_metadata,
102
- proposal_lens)
103
-
104
- assert filtered_groups == []
105
- assert indices == []
106
-
107
-
108
-def test_all_non_zero_with_zero_filter(fake_sequence_group_metadata):
109
- proposal_lens = [1, 1, 1]
110
- _, (filtered_groups,
111
- indices) = split_batch_by_proposal_len(fake_sequence_group_metadata,
112
- proposal_lens)
113
-
114
- assert filtered_groups == []
115
- assert indices == []
116
-
117
-
118
-def mock_spec_decode_sampler(acceptance_sampler_method):
119
- """
120
- Returns either a RejectionSampler or TypicalAcceptanceSampler
121
- object depending on whether acceptance_sampler_method is
122
- 'rejection_sampler' or 'typical_acceptance_sampler' respectively.
123
- """
124
- if acceptance_sampler_method == "rejection_sampler":
125
- sampler = MagicMock(spec=RejectionSampler)
126
- sampler.token_id_dtype = torch.int64
127
- return sampler
128
- elif acceptance_sampler_method == "typical_acceptance_sampler":
129
- sampler = MagicMock(spec=TypicalAcceptanceSampler)
130
- sampler.token_id_dtype = torch.int64
131
- return sampler
132
- else:
133
- raise ValueError(f"Invalid sampler name {acceptance_sampler_method}")
134
-
135
-
136
-def test_get_sampled_token_logprobs():
137
- """Verify get_sampled_token_logprobs returns consistent rankings
138
- with regular get_ranks when probabilities match exactly.
139
- """
140
- logprob_tensor = torch.tensor(
141
- [[[-.1, -.1]] * 2]) # shape (num_steps, batch_size, vocab_size)
142
- sampled_token_tensor = torch.tensor([[1,
143
- 0]]) # shape (num_steps, batch_size)
144
- ranks_spec_dec, _ = get_sampled_token_logprobs(logprob_tensor,
145
- sampled_token_tensor)
146
-
147
- ranks_regular = _get_ranks(logprob_tensor.reshape((2, -1)),
148
- sampled_token_tensor.reshape(-1))
149
-
150
- assert torch.equal(ranks_spec_dec.reshape(-1), ranks_regular)
tests/spec_decode/utils.py
deleted
-290
@@ -1,290 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from collections.abc import Sequence as GenericSequence
5
-from itertools import count
6
-from typing import Callable, Optional, TypeVar, Union
7
-from unittest.mock import MagicMock
8
-
9
-import torch
10
-
11
-from vllm.engine.arg_utils import EngineArgs
12
-from vllm.model_executor.layers.sampler import SamplerOutput
13
-from vllm.model_executor.utils import set_random_seed
14
-from vllm.sampling_params import SamplingParams
15
-from vllm.sequence import (CompletionSequenceGroupOutput, Logprob,
16
- SequenceData, SequenceGroupMetadata, SequenceOutput)
17
-from vllm.utils import get_distributed_init_method, get_ip, get_open_port
18
-from vllm.worker.cache_engine import CacheEngine
19
-from vllm.worker.model_runner import ModelRunner
20
-from vllm.worker.worker import Worker
21
-
22
-T = TypeVar("T", bound=Worker)
23
-
24
-
25
-def round_up_to_next_block(seq_len: int, block_size: int) -> int:
26
- return (seq_len + block_size - 1) // block_size
27
-
28
-
29
-def mock_worker(cls=None,
30
- vocab_size: int = 30_000,
31
- max_model_len: int = 2048,
32
- rank: int = 0,
33
- use_spec: bool = True) -> MagicMock:
34
- if cls is None:
35
- cls = Worker
36
-
37
- spec = cls if use_spec else None
38
-
39
- worker = MagicMock(spec=spec)
40
- worker.vocab_size = vocab_size
41
- worker.max_model_len = max_model_len
42
- worker.rank = rank
43
- worker.device = 'cuda:0'
44
- return worker
45
-
46
-
47
-def patch_execute_model_with_seeds(worker: Worker, rand_seeds: list[int]):
48
- seed_iter = iter(rand_seeds)
49
- original_execute_model = worker.execute_model
50
-
51
- def new_execute_model(*args, **kwargs):
52
- result = original_execute_model(*args, **kwargs)
53
- set_random_seed(next(seed_iter))
54
- return result
55
-
56
- return new_execute_model
57
-
58
-
59
-def zero_kv_cache(cache_engine: list[CacheEngine]):
60
- assert cache_engine[0].gpu_cache
61
- for key_blocks, value_blocks in cache_engine[0].gpu_cache:
62
- key_blocks.zero_()
63
- value_blocks.zero_()
64
-
65
-
66
-def create_worker(cls: Callable[..., T],
67
- model_name: str,
68
- block_size: int,
69
- num_gpu_blocks: int,
70
- seed: int,
71
- is_driver_worker: bool = True,
72
- enforce_eager: bool = True,
73
- model_runner_cls: Optional[ModelRunner] = None,
74
- dtype: Optional[str] = "auto") -> T:
75
- engine_args = EngineArgs(
76
- model=model_name,
77
- seed=seed,
78
- block_size=block_size,
79
- enforce_eager=enforce_eager,
80
- dtype=dtype,
81
- )
82
- engine_config = engine_args.create_engine_config()
83
-
84
- distributed_init_method = get_distributed_init_method(
85
- get_ip(), get_open_port())
86
-
87
- worker = cls(
88
- vllm_config=engine_config,
89
- local_rank=0,
90
- rank=0,
91
- distributed_init_method=distributed_init_method,
92
- is_driver_worker=is_driver_worker,
93
- model_runner_cls=model_runner_cls,
94
- )
95
-
96
- worker.init_device()
97
- worker.load_model()
98
-
99
- engine_config.cache_config.num_gpu_blocks = num_gpu_blocks
100
- engine_config.cache_config.num_cpu_blocks = 0
101
- worker.initialize_cache(
102
- num_gpu_blocks=engine_config.cache_config.num_gpu_blocks,
103
- num_cpu_blocks=engine_config.cache_config.num_cpu_blocks)
104
-
105
- return worker
106
-
107
-
108
-def create_seq_group_metadata_from_prompts(
109
- prompts: list[list[int]],
110
- num_gpu_blocks: int,
111
- block_size: int,
112
- final_prompt_lens: list[int],
113
- continuations: Optional[list[list[int]]] = None,
114
- seq_ids: Optional[list[int]] = None,
115
-) -> list[SequenceGroupMetadata]:
116
-
117
- if continuations is None:
118
- continuations = [[] for _ in prompts]
119
-
120
- if seq_ids is None:
121
- seq_ids = list(i for i, _ in enumerate(prompts))
122
-
123
- free_gpu_blocks = list(range(num_gpu_blocks))
124
-
125
- block_allocations = {
126
- i: [
127
- free_gpu_blocks.pop()
128
- for _ in range(round_up_to_next_block(final_len, block_size))
129
- ]
130
- for i, final_len in enumerate(final_prompt_lens)
131
- }
132
-
133
- seq_grou_metadata_list = []
134
- for i, (prompt_token_ids,
135
- cont_token_ids) in enumerate(zip(prompts, continuations)):
136
- data = SequenceData.from_seqs(prompt_token_ids, cont_token_ids)
137
- data.update_num_computed_tokens(
138
- len(prompt_token_ids) + len(cont_token_ids) - 1)
139
- seq_data = {i: data}
140
- seq_grou_metadata_list.append(
141
- SequenceGroupMetadata(
142
- request_id=str(i),
143
- is_prompt=len(cont_token_ids) == 0,
144
- seq_data=seq_data,
145
- sampling_params=SamplingParams(temperature=0.0),
146
- block_tables={i: block_allocations[i][:]},
147
- ))
148
- return seq_grou_metadata_list
149
-
150
-
151
-def create_chunked_seq_group_metadata_from_prompt(
152
- prompt: list[int],
153
- num_gpu_blocks: int,
154
- chunk_size: int,
155
- block_size: int,
156
- seq_id: Optional[int] = None) -> list[SequenceGroupMetadata]:
157
-
158
- if seq_id is None:
159
- seq_id = 0
160
-
161
- free_gpu_blocks = list(range(num_gpu_blocks))
162
-
163
- block_allocations = [
164
- free_gpu_blocks.pop()
165
- for _ in range(round_up_to_next_block(len(prompt), block_size))
166
- ]
167
-
168
- seq_group_metadata_list = []
169
- for i, idx in enumerate(range(0, len(prompt), chunk_size)):
170
- chunk_ids = prompt[idx:idx + chunk_size]
171
- data = SequenceData.from_seqs(prompt)
172
- data.update_num_computed_tokens(idx)
173
- seq_data = {i: data}
174
- seq_group_metadata_list.append(
175
- SequenceGroupMetadata(
176
- request_id=str(seq_id),
177
- is_prompt=True,
178
- do_sample=idx + chunk_size >= len(prompt), # terminal chunk
179
- seq_data=seq_data,
180
- sampling_params=SamplingParams(temperature=0.0),
181
- block_tables={i: block_allocations},
182
- token_chunk_size=len(chunk_ids)))
183
- return seq_group_metadata_list
184
-
185
-
186
-def assert_logprobs_dict_allclose(
187
- actual_logprobs: list[dict[int, Logprob]],
188
- expected_logprobs: list[dict[int, Logprob]]) -> None:
189
- for single_step_actual_logprobs, single_step_expected_logprobs in zip(
190
- actual_logprobs, expected_logprobs):
191
- assert set(single_step_actual_logprobs.keys()) == set(
192
- single_step_expected_logprobs.keys())
193
- for token_id in single_step_actual_logprobs:
194
- actual = torch.tensor(
195
- single_step_actual_logprobs[token_id].logprob)
196
- expected = torch.tensor(
197
- single_step_expected_logprobs[token_id].logprob)
198
- torch.testing.assert_close(actual, expected)
199
-
200
-
201
-def create_sampler_output_list(
202
- token_ids: torch.Tensor,
203
- probs: GenericSequence[Optional[torch.Tensor]],
204
- logprobs: GenericSequence[Optional[torch.Tensor]],
205
- seq_ids: Optional[list[int]] = None) -> list[SamplerOutput]:
206
- num_steps, batch_size = token_ids.shape
207
- token_ids_by_step = token_ids.tolist()
208
-
209
- if seq_ids is None:
210
- seq_ids = list(range(batch_size))
211
-
212
- return [
213
- SamplerOutput(outputs=[
214
- CompletionSequenceGroupOutput(
215
- samples=[
216
- SequenceOutput(
217
- output_token=token_id,
218
- parent_seq_id=seq_ids[seq_index],
219
- logprobs={token_id: Logprob(0)},
220
- )
221
- ],
222
- prompt_logprobs=None,
223
- ) for seq_index, token_id in enumerate(token_ids_by_step[step])
224
- ],
225
- sampled_token_probs=probs[step],
226
- logprobs=logprobs[step],
227
- sampled_token_ids=token_ids[step])
228
- for step in range(num_steps)
229
- ]
230
-
231
-
232
-def create_batch(batch_size,
233
- k,
234
- prompt_len: Union[int, list[int]] = 10,
235
- prev_output_token_len: int = 10,
236
- seq_ids: Optional[list[int]] = None,
237
- num_gpu_blocks: Optional[int] = None,
238
- block_size: Optional[int] = None,
239
- prefill_chunk_size: Optional[int] = None):
240
- if block_size is None:
241
- block_size = 8
242
-
243
- if num_gpu_blocks is None:
244
- num_gpu_blocks = 2048 // block_size
245
-
246
- iterator = count()
247
-
248
- if isinstance(prompt_len, int):
249
- prompt_lens = [prompt_len for _ in range(batch_size)]
250
- else:
251
- prompt_lens = prompt_len
252
-
253
- prompts = [[next(iterator) for _ in range(p_len)] for p_len in prompt_lens]
254
-
255
- if prefill_chunk_size:
256
- # Create a batch of chunked prompts.
257
- if not seq_ids:
258
- seq_ids = list(range(len(prompts)))
259
- seq_group_metadata_list = []
260
- for p, sid in zip(prompts, seq_ids):
261
- seq_group_metadata_list += \
262
- create_chunked_seq_group_metadata_from_prompt(
263
- p, num_gpu_blocks, prefill_chunk_size, block_size, sid)
264
- seq_group_metadata_list = seq_group_metadata_list[:batch_size]
265
- prev_output_tokens = []
266
- else:
267
- prev_output_tokens = [[
268
- next(iterator) for _ in range(prev_output_token_len)
269
- ] for _ in range(batch_size)]
270
- final_prompt_lens = [
271
- len(prompt) + len(prev_output_token) + k + 1
272
- for prompt, prev_output_token in zip(prompts, prev_output_tokens)
273
- ]
274
-
275
- seq_group_metadata_list = create_seq_group_metadata_from_prompts(
276
- prompts, num_gpu_blocks, block_size, final_prompt_lens,
277
- prev_output_tokens, seq_ids)
278
- return seq_group_metadata_list, prompts, prev_output_tokens
279
-
280
-
281
-def maybe_enable_chunked_prefill(prefill_chunk_size, llm_kwargs):
282
- if prefill_chunk_size > 0:
283
- llm_kwargs.update(
284
- **{
285
- "enable_chunked_prefill": True,
286
- "max_num_batched_tokens": prefill_chunk_size,
287
- "max_num_seqs": prefill_chunk_size
288
- })
289
- else:
290
- llm_kwargs["enable_chunked_prefill"] = False
tests/test_sequence.py
-1
@@ -29,7 +29,6 @@ def test_sampler_output_initialization(sampler_output, sample_outputs):
29
assert len(sampler_output) == len(sample_outputs)
30
assert sampler_output.sampled_token_probs is None
31
assert sampler_output.sampled_token_ids is None
32
- assert sampler_output.spec_decode_worker_metrics is None
32
33
34
def test_sampler_output_getitem(sampler_output, sample_outputs):
tests/v1/test_oracle.py
-6
@@ -40,12 +40,6 @@ def test_unsupported_configs(monkeypatch):
40
with monkeypatch.context() as m:
41
m.setenv("VLLM_USE_V1", "1")
42
43
- with pytest.raises(NotImplementedError):
44
- AsyncEngineArgs(
45
- model=MODEL,
46
- kv_cache_dtype="fp8",
47
- ).create_engine_config()
48
-
43
with pytest.raises(NotImplementedError):
44
AsyncEngineArgs(
45
model=MODEL,
tools/mypy.sh
-1
@@ -32,6 +32,5 @@ run_mypy vllm/lora
32
run_mypy vllm/model_executor
33
run_mypy vllm/plugins
34
run_mypy vllm/prompt_adapter
35
-run_mypy vllm/spec_decode
35
run_mypy vllm/worker
36
run_mypy vllm/v1
vllm/config.py
+7
-54
@@ -2536,8 +2536,6 @@ class DeviceConfig:
2536
2537
SpeculativeMethod = Literal["ngram", "eagle", "eagle3", "medusa",
2538
"mlp_speculator", "draft_model", "deepseek_mtp"]
2539
-SpeculativeAcceptanceMethod = Literal["rejection_sampler",
2540
- "typical_acceptance_sampler"]
2539
2540
2541
@config
@@ -2560,13 +2558,6 @@ class SpeculativeConfig:
2558
2559
If using `ngram` method, the related configuration `prompt_lookup_max` and
2560
`prompt_lookup_min` should be considered."""
2563
- acceptance_method: SpeculativeAcceptanceMethod = "rejection_sampler"
2564
- """The method to use for accepting draft tokens:\n
2565
- - "rejection_sampler" maps to `RejectionSampler`.\n
2566
- - "typical_acceptance_sampler" maps to `TypicalAcceptanceSampler`.
2567
-
2568
- If using `typical_acceptance_sampler`, the related configuration
2569
- `posterior_threshold` and `posterior_alpha` should be considered."""
2561
draft_tensor_parallel_size: Optional[int] = None
2562
"""The degree of the tensor parallelism for the draft model. Can only be 1
2563
or the same as the target model's tensor parallel size."""
@@ -2593,9 +2584,6 @@ class SpeculativeConfig:
2584
will use the default version."""
2585
2586
# Advanced control
2596
- disable_mqa_scorer: bool = False
2597
- """Disable the MQA scorer and fall back to batch expansion for scoring
2598
- proposals."""
2587
disable_by_batch_size: Optional[int] = None
2588
"""Disable speculative decoding for new incoming requests when the number
2589
of enqueued requests is larger than this value, if provided."""
@@ -2608,16 +2596,6 @@ class SpeculativeConfig:
2596
"""Minimum size of ngram token window when using Ngram proposer, if
2597
provided. Defaults to 1."""
2598
2611
- # Typical acceptance sampler configuration
2612
- posterior_threshold: Optional[float] = None
2613
- """A threshold value that sets a lower bound on the posterior probability
2614
- of a token in the target model for it to be accepted. This threshold is
2615
- used only when we use the `TypicalAcceptanceSampler` for token acceptance.
2616
- """
2617
- posterior_alpha: Optional[float] = None
2618
- """Scaling factor for entropy-based threshold, applied when using
2619
- `TypicalAcceptanceSampler`."""
2620
-
2599
speculative_token_tree: Optional[str] = None
2600
"""Specifies the tree structure for speculative token generation.
2601
"""
@@ -2795,8 +2773,8 @@ class SpeculativeConfig:
2773
elif (self.draft_model_config.hf_config.model_type ==
2774
"mlp_speculator"):
2775
self.method = "mlp_speculator"
2798
- elif (self.draft_model_config.hf_config.model_type ==
2799
- "deepseek_mtp"):
2776
+ elif (self.draft_model_config.hf_config.model_type
2777
+ in ("deepseek_mtp", "mimo_mtp")):
2778
self.method = "deepseek_mtp"
2779
if self.num_speculative_tokens > 1:
2780
logger.warning(
@@ -2806,6 +2784,11 @@ class SpeculativeConfig:
2784
)
2785
else:
2786
self.method = "draft_model"
2787
+ raise NotImplementedError(
2788
+ "Speculative decoding with draft model is not "
2789
+ "supported yet. Please consider using other "
2790
+ "speculative decoding methods such as ngram, medusa, "
2791
+ "eagle, or deepseek_mtp.")
2792
2793
# Replace hf_config for EAGLE draft_model
2794
if self.method in ("eagle", "eagle3"):
@@ -2864,12 +2847,6 @@ class SpeculativeConfig:
2847
self.target_parallel_config,
2848
self.draft_tensor_parallel_size))
2849
2867
- if self.acceptance_method == "typical_acceptance_sampler":
2868
- if self.posterior_threshold is None:
2869
- self.posterior_threshold = 0.09
2870
- if self.posterior_alpha is None:
2871
- self.posterior_alpha = 0.3
2872
-
2850
@staticmethod
2851
def _maybe_override_draft_max_model_len(
2852
speculative_max_model_len: Optional[int],
@@ -2975,30 +2952,6 @@ class SpeculativeConfig:
2952
if self.draft_model_config:
2953
self.draft_model_config.verify_with_parallel_config(
2954
self.draft_parallel_config)
2978
- # Validate and set draft token acceptance related settings.
2979
-
2980
- if self.acceptance_method is None:
2981
- raise ValueError("acceptance_method is not set. "
2982
- "Expected values are rejection_sampler or "
2983
- "typical_acceptance_sampler.")
2984
-
2985
- if (self.acceptance_method != 'rejection_sampler'
2986
- and self.acceptance_method != 'typical_acceptance_sampler'):
2987
- raise ValueError(
2988
- "Expected acceptance_method to be either "
2989
- "rejection_sampler or typical_acceptance_sampler. Instead it "
2990
- f"is {self.acceptance_method}")
2991
-
2992
- if self.acceptance_method == "typical_acceptance_sampler" and (
2993
- (self.posterior_threshold is not None
2994
- and self.posterior_threshold < 0) or
2995
- (self.posterior_alpha is not None and self.posterior_alpha < 0)):
2996
- raise ValueError(
2997
- "Expected the posterior_threshold and posterior_alpha of "
2998
- "typical_acceptance_sampler to be > 0. "
2999
- "Instead found posterior_threshold = "
3000
- f"{self.posterior_threshold} and posterior_alpha = "
3001
- f"{self.posterior_alpha}")
2955
2956
if (self.disable_by_batch_size is not None
2957
and self.disable_by_batch_size < 2):
vllm/engine/arg_utils.py
+6
-22
@@ -1417,28 +1417,12 @@ class EngineArgs:
1417
return False
1418
1419
# V1 supports N-gram, Medusa, and Eagle speculative decoding.
1420
- is_ngram_enabled = False
1421
- is_eagle_enabled = False
1422
- is_medusa_enabled = False
1423
- if self.speculative_config is not None:
1424
- # This is supported but experimental (handled below).
1425
- speculative_method = self.speculative_config.get("method")
1426
- if speculative_method:
1427
- if speculative_method in ("ngram", "[ngram]"):
1428
- is_ngram_enabled = True
1429
- elif speculative_method == "medusa":
1430
- is_medusa_enabled = True
1431
- elif speculative_method in ("eagle", "eagle3", "deepseek_mtp"):
1432
- is_eagle_enabled = True
1433
- else:
1434
- speculative_model = self.speculative_config.get("model")
1435
- if speculative_model in ("ngram", "[ngram]"):
1436
- is_ngram_enabled = True
1437
- if not (is_ngram_enabled or is_eagle_enabled or is_medusa_enabled):
1438
- # Other speculative decoding methods are not supported yet.
1439
- _raise_or_fallback(feature_name="Speculative Decoding",
1440
- recommend_to_remove=False)
1441
- return False
1420
+ if (self.speculative_config is not None
1421
+ and self.speculative_config.get("method") == "draft_model"):
1422
+ raise NotImplementedError(
1423
+ "Speculative decoding with draft model is not supported yet. "
1424
+ "Please consider using other speculative decoding methods "
1425
+ "such as ngram, medusa, eagle, or deepseek_mtp.")
1426
1427
# No XFormers so far.
1428
V1_BACKENDS = [
vllm/engine/llm_engine.py
-8
@@ -1780,13 +1780,6 @@ class LLMEngine:
1780
num_generation_tokens_from_prefill_groups)
1781
num_tokens_iter = (num_generation_tokens_iter +
1782
num_prompt_tokens_iter)
1783
- # Spec decode, if enabled, emits specialized metrics from the worker in
1784
- # sampler output.
1785
- if model_output and isinstance(model_output[0], SamplerOutput) and (
1786
- model_output[0].spec_decode_worker_metrics is not None):
1787
- spec_decode_metrics = model_output[0].spec_decode_worker_metrics
1788
- else:
1789
- spec_decode_metrics = None
1783
1784
return Stats(
1785
now=now,
@@ -1808,7 +1801,6 @@ class LLMEngine:
1801
num_tokens_iter=num_tokens_iter,
1802
time_to_first_tokens_iter=time_to_first_tokens_iter,
1803
time_per_output_tokens_iter=time_per_output_tokens_iter,
1811
- spec_decode_metrics=spec_decode_metrics,
1804
num_preemption_iter=num_preemption_iter,
1805
1806
# Request stats
vllm/engine/metrics.py
-66
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import time
5
-from typing import TYPE_CHECKING
5
from typing import Counter as CollectionsCounter
6
from typing import Dict, List, Optional, Type, Union, cast
7
@@ -19,9 +18,6 @@ if ray is not None:
18
else:
19
ray_metrics = None
20
22
-if TYPE_CHECKING:
23
- from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics
24
-
21
logger = init_logger(__name__)
22
23
prometheus_client.disable_created_metrics()
@@ -199,30 +195,6 @@ class Metrics:
195
documentation="Count of successfully processed requests.",
196
labelnames=labelnames + [Metrics.labelname_finish_reason])
197
202
- # Speculative decoding stats
203
- self.gauge_spec_decode_draft_acceptance_rate = self._gauge_cls(
204
- name="vllm:spec_decode_draft_acceptance_rate",
205
- documentation="Speulative token acceptance rate.",
206
- labelnames=labelnames,
207
- multiprocess_mode="sum")
208
- self.gauge_spec_decode_efficiency = self._gauge_cls(
209
- name="vllm:spec_decode_efficiency",
210
- documentation="Speculative decoding system efficiency.",
211
- labelnames=labelnames,
212
- multiprocess_mode="sum")
213
- self.counter_spec_decode_num_accepted_tokens = (self._counter_cls(
214
- name="vllm:spec_decode_num_accepted_tokens_total",
215
- documentation="Number of accepted tokens.",
216
- labelnames=labelnames))
217
- self.counter_spec_decode_num_draft_tokens = self._counter_cls(
218
- name="vllm:spec_decode_num_draft_tokens_total",
219
- documentation="Number of draft tokens.",
220
- labelnames=labelnames)
221
- self.counter_spec_decode_num_emitted_tokens = (self._counter_cls(
222
- name="vllm:spec_decode_num_emitted_tokens_total",
223
- documentation="Number of emitted tokens.",
224
- labelnames=labelnames))
225
-
198
199
# --8<-- [end:metrics-definitions]
200
@@ -391,9 +363,6 @@ class LoggingStatLogger(StatLoggerBase):
363
self.num_prompt_tokens.append(stats.num_prompt_tokens_iter)
364
self.num_generation_tokens.append(stats.num_generation_tokens_iter)
365
394
- # Update spec decode metrics
395
- self.maybe_update_spec_decode_metrics(stats)
396
-
366
# Log locally every local_interval seconds.
367
if local_interval_elapsed(stats.now, self.last_local_log,
368
self.local_interval):
@@ -435,10 +404,6 @@ class LoggingStatLogger(StatLoggerBase):
404
stats.gpu_prefix_cache_hit_rate * 100,
405
stats.cpu_prefix_cache_hit_rate * 100,
406
)
438
- if self.spec_decode_metrics is not None:
439
- log_fn(
440
- self._format_spec_decode_metrics_str(
441
- self.spec_decode_metrics))
407
408
self._reset(stats, prompt_throughput, generation_throughput)
409
@@ -447,21 +412,9 @@ class LoggingStatLogger(StatLoggerBase):
412
self.num_prompt_tokens = []
413
self.num_generation_tokens = []
414
self.last_local_log = stats.now
450
- self.spec_decode_metrics = None
415
self.last_prompt_throughput = prompt_throughput
416
self.last_generation_throughput = generation_throughput
417
454
- def _format_spec_decode_metrics_str(
455
- self, metrics: "SpecDecodeWorkerMetrics") -> str:
456
-
457
- return ("Speculative metrics: "
458
- f"Draft acceptance rate: {metrics.draft_acceptance_rate:.3f}, "
459
- f"System efficiency: {metrics.system_efficiency:.3f}, "
460
- f"Number of speculative tokens: {metrics.num_spec_tokens}, "
461
- f"Number of accepted tokens: {metrics.accepted_tokens}, "
462
- f"Number of draft tokens: {metrics.draft_tokens}, "
463
- f"Number of emitted tokens: {metrics.emitted_tokens}.")
464
-
418
def info(self, type: str, obj: SupportsMetricsInfo) -> None:
419
raise NotImplementedError
420
@@ -579,33 +532,14 @@ class PrometheusStatLogger(StatLoggerBase):
532
self.num_prompt_tokens.append(stats.num_prompt_tokens_iter)
533
self.num_generation_tokens.append(stats.num_generation_tokens_iter)
534
582
- # Update spec decode metrics
583
- self.maybe_update_spec_decode_metrics(stats)
584
-
535
# Log locally every local_interval seconds.
536
if local_interval_elapsed(stats.now, self.last_local_log,
537
self.local_interval):
588
- if self.spec_decode_metrics is not None:
589
- self._log_gauge(
590
- self.metrics.gauge_spec_decode_draft_acceptance_rate,
591
- self.spec_decode_metrics.draft_acceptance_rate)
592
- self._log_gauge(self.metrics.gauge_spec_decode_efficiency,
593
- self.spec_decode_metrics.system_efficiency)
594
- self._log_counter(
595
- self.metrics.counter_spec_decode_num_accepted_tokens,
596
- self.spec_decode_metrics.accepted_tokens)
597
- self._log_counter(
598
- self.metrics.counter_spec_decode_num_draft_tokens,
599
- self.spec_decode_metrics.draft_tokens)
600
- self._log_counter(
601
- self.metrics.counter_spec_decode_num_emitted_tokens,
602
- self.spec_decode_metrics.emitted_tokens)
538
539
# Reset tracked stats for next interval.
540
self.num_prompt_tokens = []
541
self.num_generation_tokens = []
542
self.last_local_log = stats.now
608
- self.spec_decode_metrics = None
543
544
def info(self, type: str, obj: SupportsMetricsInfo) -> None:
545
# Info type metrics are syntactic sugar for a gauge permanently set to 1
vllm/engine/metrics_types.py
+1
-11
@@ -16,10 +16,9 @@ do this in Python code and lazily import prometheus_client.
16
import time
17
from abc import ABC, abstractmethod
18
from dataclasses import dataclass
19
-from typing import List, Optional
19
+from typing import List
20
21
from vllm.config import SupportsMetricsInfo, VllmConfig
22
-from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics
22
23
24
@dataclass
@@ -65,8 +64,6 @@ class Stats:
64
running_lora_adapters: List[str]
65
max_lora: str
66
68
- spec_decode_metrics: Optional["SpecDecodeWorkerMetrics"] = None
69
-
67
68
class StatLoggerBase(ABC):
69
"""Base class for StatLogger."""
@@ -77,7 +74,6 @@ class StatLoggerBase(ABC):
74
self.num_generation_tokens: List[int] = []
75
self.last_local_log = time.time()
76
self.local_interval = local_interval
80
- self.spec_decode_metrics: Optional[SpecDecodeWorkerMetrics] = None
77
78
@abstractmethod
79
def log(self, stats: Stats) -> None:
@@ -86,9 +82,3 @@ class StatLoggerBase(ABC):
82
@abstractmethod
83
def info(self, type: str, obj: SupportsMetricsInfo) -> None:
84
raise NotImplementedError
89
-
90
- def maybe_update_spec_decode_metrics(self, stats: Stats):
91
- """Save spec decode metrics (since they are unlikely
92
- to be emitted at same time as log interval)."""
93
- if stats.spec_decode_metrics is not None:
94
- self.spec_decode_metrics = stats.spec_decode_metrics
vllm/engine/output_processor/multi_step.py
-5
@@ -104,11 +104,6 @@ class MultiStepOutputProcessor(SequenceGroupOutputProcessor):
104
seqs = sequence_group.get_seqs(
105
status=SequenceStatus.FINISHED_ABORTED)
106
107
- for output in outputs:
108
- if output.samples[0].output_token != VLLM_INVALID_TOKEN_ID:
109
- sequence_group.metrics.spec_token_acceptance_counts[
110
- output.step_index] += 1
111
-
107
assert seqs, "Expected RUNNING or FINISHED_ABORTED sequences"
108
assert len(seqs) == 1, (
109
"Beam search not supported in multi-step decoding.")
vllm/model_executor/layers/rejection_sampler.py
deleted
-406
@@ -1,406 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from functools import cached_property
5
-from importlib.util import find_spec
6
-from typing import Optional
7
-
8
-import torch
9
-import torch.jit
10
-
11
-import vllm.envs as envs
12
-from vllm.logger import init_logger
13
-from vllm.model_executor.layers.spec_decode_base_sampler import (
14
- SpecDecodeStochasticBaseSampler)
15
-from vllm.platforms import current_platform
16
-
17
-logger = init_logger(__name__)
18
-
19
-if find_spec("flashinfer"):
20
- """
21
- Consider utilizing the FlashInfer rejection sampling kernel initially,
22
- as it employs a dedicated kernel rather than relying on
23
- Torch tensor operations. This design choice helps to fuse operations,
24
- reduce memory I/O, and consequently enhances performance.
25
- """
26
- from flashinfer.sampling import chain_speculative_sampling
27
-else:
28
- chain_speculative_sampling = None
29
-
30
-
31
-class RejectionSampler(SpecDecodeStochasticBaseSampler):
32
- """Apply modified rejection sampling as described in "Accelerating Large
33
- Language Model Decoding with Speculative Sampling"
34
- https://arxiv.org/pdf/2302.01318.pdf.
35
- """
36
-
37
- def __init__(self,
38
- strict_mode: bool = False,
39
- use_flashinfer: Optional[bool] = None):
40
- """Create a rejection sampler.
41
-
42
- Args:
43
- strict_mode: Whether or not to perform shape/device/dtype checks
44
- during sampling. This catches correctness issues but adds
45
- nontrivial latency.
46
- use_flashinfer: We will use this parameter to determine whether
47
- to use the FlashInfer rejection sampling kernel or not. If it's
48
- None, we will use the default value from the environment variable.
49
- This parameter is only used for testing purposes.
50
- """
51
- super().__init__(strict_mode=strict_mode)
52
- if use_flashinfer is None:
53
- self.use_flashinfer = envs.VLLM_USE_FLASHINFER_SAMPLER and (
54
- chain_speculative_sampling is not None)
55
- else:
56
- self.use_flashinfer = use_flashinfer
57
-
58
- if self.use_flashinfer:
59
- logger.info("Use flashinfer for rejection sampling.")
60
- else:
61
- logger.info("Use pytorch for rejection sampling.")
62
-
63
- def forward(
64
- self,
65
- target_with_bonus_probs: torch.Tensor,
66
- bonus_token_ids: torch.Tensor,
67
- draft_probs: torch.Tensor,
68
- draft_token_ids: torch.Tensor,
69
- seeded_seqs: Optional[dict[int, torch.Generator]] = None,
70
- ) -> torch.Tensor:
71
- """Sample token ids using rejection sampling. This accepts or rejects
72
- tokens proposed by the draft model using the probability of each token
73
- according to the draft and target models.
74
-
75
- In the worst case where all draft tokens are rejected, it is guaranteed
76
- one correct token will be emitted.
77
-
78
- In the case where all draft tokens are accepted, a bonus token will be
79
- accepted as its cheap to have the target model score this speculative
80
- sequence.
81
-
82
- Args:
83
- target_with_bonus_probs: The probability distribution
84
- over token ids given context according to the target model.
85
- shape = [batch_size, num_speculative_tokens + 1, vocab_size]
86
-
87
- bonus_token_ids: The "bonus" token ids that are accepted iff all
88
- speculative tokens in a sequence are accepted.
89
- shape = [batch_size, num_bonus_tokens]
90
-
91
- draft_probs: The probability distribution over token ids given
92
- context according to the draft model.
93
- shape = [batch_size, num_speculative_tokens, vocab_size]
94
-
95
- draft_token_ids: The token ids that were sampled from the draft
96
- probabilities.
97
- shape = [batch_size, num_speculative_tokens]
98
-
99
- seeded_seqs: Dict of batch row index to torch generator, for
100
- sequences using seeded generation.
101
-
102
- Returns:
103
- output_token_ids: The token ids sampled via rejection sampling,
104
- or -1 if unable to sample a token because the previous token
105
- was rejected.
106
- shape = [batch_size, num_speculative_tokens + num_bonus_tokens]
107
- """
108
- # Only perform shape/dtype/device checking in strict mode, as it adds
109
- # overhead.
110
- if self._strict_mode:
111
- self._raise_if_incorrect_input(target_with_bonus_probs,
112
- draft_token_ids, bonus_token_ids,
113
- draft_probs)
114
-
115
- batch_size, k, _ = draft_probs.shape
116
-
117
- # batch_size = 0 when all requests in the batch are
118
- # non_spec requests. In this case, output_token_ids is
119
- # just an empty tensor.
120
- if batch_size == 0:
121
- return torch.empty(0, k + 1, device=draft_probs.device, dtype=int)
122
-
123
- # If use Flashinfer chain_speculative_sampling kernel
124
- # for rejection sampling
125
- if self.use_flashinfer and chain_speculative_sampling is not None:
126
- batch_size, k, _ = draft_probs.shape
127
-
128
- (output_token_ids, accepted_token_num,
129
- emitted_token_num) = chain_speculative_sampling(
130
- draft_probs,
131
- draft_token_ids,
132
- target_with_bonus_probs,
133
- )
134
-
135
- # num_emitted_tokens returned by flashinfer
136
- # does not include the bonus token
137
- # Flashinfer stops at the first token that violates
138
- # the condition p >= q and does not include recovery/bonus token.
139
- # Therefore, we need to add batch_size here.
140
- self.num_accepted_tokens += accepted_token_num.sum()
141
- self.num_emitted_tokens += emitted_token_num.sum() + batch_size
142
- self.num_draft_tokens += batch_size * k
143
- else:
144
- accepted, recovered_token_ids = (
145
- self._batch_modified_rejection_sampling(
146
- target_with_bonus_probs[:, :-1],
147
- draft_probs,
148
- draft_token_ids,
149
- seeded_seqs,
150
- ))
151
-
152
- output_token_ids = self._create_output(
153
- accepted,
154
- recovered_token_ids,
155
- draft_token_ids,
156
- bonus_token_ids,
157
- )
158
-
159
- return output_token_ids
160
-
161
- def _batch_modified_rejection_sampling(
162
- self,
163
- target_probs: torch.Tensor, # [batch_size, k, vocab_size]
164
- draft_probs: torch.Tensor, # [batch_size, k, vocab_size]
165
- draft_token_ids: torch.Tensor, # [batch_size, k]
166
- seeded_seqs: Optional[dict[int, torch.Generator]],
167
- ) -> tuple[torch.Tensor, torch.Tensor]:
168
- """Perform modified rejection sampling on each sequence.
169
-
170
- Returns:
171
- A tuple of two tensors:
172
- 0: A bool tensor of which tokens in each sequence is accepted.
173
- shape = [batch_size, k]
174
- 1: Token ids sampled from a recovered distribution, to be used
175
- when a token is rejected.
176
- shape = [batch_size, k]
177
- """
178
-
179
- batch_size, k, vocab_size = draft_probs.shape
180
-
181
- # shape [batch_size, k]
182
- accepted = self._get_accepted(target_probs, draft_probs,
183
- draft_token_ids, seeded_seqs)
184
-
185
- recovered_probs = self._get_recovered_probs(
186
- target_probs, draft_probs).reshape(batch_size * k, vocab_size)
187
-
188
- # NOTE: the recovered_probs are overwritten by this method.
189
- recovered_token_ids = _multinomial(
190
- recovered_probs,
191
- num_samples=1,
192
- k=k,
193
- seeded_seqs=seeded_seqs or {},
194
- ).reshape(batch_size, k)
195
-
196
- return accepted, recovered_token_ids
197
-
198
- def _create_uniform_samples(self,
199
- seeded_seqs: Optional[dict[int,
200
- torch.Generator]],
201
- batch_size: int, k: int,
202
- device: torch.device) -> torch.Tensor:
203
- """
204
- Generates a batch of uniform random samples, with optional seeding
205
- for specific sequences.
206
-
207
- This method creates a tensor of shape `(batch_size, k + 1)` filled
208
- with uniform random values in the range [0, 1). If `seeded_seqs`
209
- is provided, the sequences corresponding to specific indices
210
- will be generated using the provided `torch.Generator` for
211
- reproducibility. The other sequences will be generated without
212
- a seed.
213
-
214
- Args:
215
- seeded_seqs : Optional[dict[int, torch.Generator]]
216
- A dictionary mapping indices in the batch to
217
- `torch.Generator` objects. If `None`, all samples are
218
- generated without a seed.
219
- batch_size : int
220
- The number of sequences to generate.
221
- k : int
222
- The number of random samples per sequence.
223
- device : torch.device
224
- The device on which to allocate the tensor.
225
-
226
- Returns:
227
- uniform_rand : torch.Tensor
228
- A tensor of shape `(batch_size, k + 1)` containing uniform
229
- random values in the range [0, 1).
230
- """
231
- if not seeded_seqs:
232
- return torch.rand(batch_size, k + 1, device=device)
233
-
234
- uniform_rand = torch.empty(batch_size, k + 1, device=device)
235
-
236
- non_seeded_indices = []
237
- for idx in range(batch_size):
238
- generator = seeded_seqs.get(idx)
239
- if generator is None:
240
- non_seeded_indices.append(idx)
241
- else:
242
- uniform_rand[idx, :] = torch.rand(1,
243
- k + 1,
244
- dtype=self.probs_dtype,
245
- device=device,
246
- generator=generator)
247
- if non_seeded_indices:
248
- uniform_rand[non_seeded_indices, :] = torch.rand(
249
- len(non_seeded_indices),
250
- k + 1,
251
- dtype=self.probs_dtype,
252
- device=device)
253
- return uniform_rand
254
-
255
- def _get_accepted(
256
- self,
257
- target_probs: torch.Tensor, # [batch_size, k, vocab_size]
258
- draft_probs: torch.Tensor, # [batch_size, k, vocab_size]
259
- draft_token_ids: torch.Tensor, # [batch_size, k]
260
- seeded_seqs: Optional[dict[int, torch.Generator]],
261
- ) -> torch.Tensor:
262
- r"""Create bool matrix over the proposed draft tokens. If
263
- True, then a token can be accepted, else it should be
264
- rejected.
265
-
266
- Given $q(\hat{x}_{n+1}|x_1, \dots, x_n)$, the probability of
267
- $\hat{x}_{n+1}$ given context $x_1, \dots, x_n$ according
268
- to the target model, and $p(\hat{x}_{n+1}|x_1, \dots, x_n)$, the
269
- same conditional probability according to the draft model, the token
270
- is accepted with probability:
271
-
272
- $$
273
- \min\left(1, \frac{q(\hat{x}_{n+1}|x_1, \dots, x_n)}
274
- {p(\hat{x}_{n+1}|x_1, \dots, x_n)}\right)
275
- $$
276
-
277
- This implementation does not apply causality. When using the output,
278
- if a token is rejected, subsequent tokens should not be used.
279
-
280
- Returns a bool tensor of shape [batch_size, k] specifying which tokens
281
- are accepted.
282
- """
283
- batch_size, k, _ = draft_probs.shape
284
- batch_indices = torch.arange(batch_size,
285
- device=target_probs.device)[:, None]
286
- probs_indices = torch.arange(k, device=target_probs.device)
287
-
288
- # shape [batch_size, k]
289
- selected_draft_probs = draft_probs[batch_indices, probs_indices,
290
- draft_token_ids]
291
-
292
- # shape [batch_size, k]
293
- selected_target_probs = target_probs[batch_indices, probs_indices,
294
- draft_token_ids]
295
-
296
- uniform_rand = self._create_uniform_samples(seeded_seqs, batch_size,
297
- k - 1, target_probs.device)
298
-
299
- capped_ratio = torch.minimum(
300
- selected_target_probs / selected_draft_probs,
301
- torch.full((1, ), 1, device=target_probs.device))
302
- accepted = uniform_rand < capped_ratio
303
-
304
- return accepted
305
-
306
- def _get_recovered_probs(
307
- self,
308
- target_probs: torch.Tensor, # [k, vocab_size]
309
- draft_probs: torch.Tensor, # [k, vocab_size]
310
- ) -> torch.Tensor:
311
- r"""Create a probability distribution for each proposed token which can
312
- be sampled if the proposed token is rejected.
313
-
314
- When this routine is applied sequentially, the true distribution of the
315
- target model is recovered (within hardware numerics).
316
-
317
- The probability distribution used in this rejection case is constructed
318
- as follows. Given $q(x|x_1, \dots, x_n)$, the probability of
319
- $x$ given context $x_1, \dots, x_n$ according to the target
320
- model and $p(x|x_1, \dots, x_n)$, the same conditional probability
321
- according to the draft model:
322
-
323
- $$
324
- x_{n+1} \sim (q(x|x_1, \dots, x_n) - p(x|x_1, \dots, x_n))_+
325
- $$
326
-
327
- where $(f(x))_+$ is defined as:
328
-
329
- $$
330
- (f(x))_+ = \frac{\max(0, f(x))}{\sum_x \max(0, f(x))}
331
- $$
332
-
333
- See https://github.com/vllm-project/vllm/pull/2336 for a visualization
334
- of the draft, target, and recovered probability distributions.
335
-
336
- Returns a tensor of shape [batch_size, k, vocab_size].
337
-
338
- Note:
339
- This batches operations on GPU and thus constructs the recovered
340
- distribution for all tokens, even if they are accepted. This causes
341
- division-by-zero errors, so we use self._smallest_positive_value to
342
- avoid that. This introduces some drift to the distribution.
343
- """
344
- _, k, _ = draft_probs.shape
345
-
346
- # shape [batch_size, k, vocab_size]
347
- difference = target_probs - draft_probs
348
-
349
- # TODO(cade): Can we use logprobs instead of probs, and avoid the
350
- # division-by-zero errors without introducing distribution drift?
351
-
352
- # shape [batch_size, k, vocab_size]
353
- f = torch.clamp(difference, min=self._smallest_positive_value)
354
-
355
- # shape [batch_size, k, vocab_size]
356
- recovered_probs = f / torch.sum(f, dim=-1).reshape(-1, k, 1)
357
-
358
- return recovered_probs
359
-
360
- @cached_property
361
- def _smallest_positive_value(self) -> float:
362
- """Return the smallest positive value representable by the probs dtype.
363
- This value is used when constructing a distribution from which to sample
364
- recovered tokens in the first rejection case.
365
-
366
- See _get_recovered_probs for more details
367
-
368
- Note that this isn't actually the smallest positive value representable
369
- by float32, but the smallest positive normal value.
370
- See https://en.wikipedia.org/wiki/Subnormal_number for more information.
371
- """
372
- return torch.finfo(self.probs_dtype).tiny
373
-
374
-
375
-# torch.multinomial forces a GPU<->CPU sync.
376
-# Therefore, we use an optimized implementation instead that skips the sync.
377
-# Note that we always sample with replacement.
378
-# probs will be modified in place, but this is fine, as we pass
379
-# in a copy already.
380
-@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
381
-def _multinomial(
382
- probs: torch.Tensor,
383
- num_samples: int,
384
- k: int,
385
- seeded_seqs: dict[int, torch.Generator],
386
-) -> torch.Tensor:
387
-
388
- if num_samples > 1:
389
- # This is equivalent to torch.repeat_interleaved (which also
390
- # forces a GPU<->CPU sync).
391
- probs = probs[:, None, :].expand(probs.shape[0], num_samples,
392
- probs.shape[1]).contiguous().view(
393
- -1, probs.shape[1])
394
- q = torch.empty_like(probs)
395
- if not seeded_seqs:
396
- q.exponential_(1.0)
397
- else:
398
- start = 0
399
- for idx in range(len(q) // k):
400
- end = start + k
401
- generator = seeded_seqs.get(idx)
402
- # Note: generator might be None for non seeded
403
- q[start:end].exponential_(1.0, generator=generator)
404
- start = end
405
-
406
- return probs.div_(q).argmax(dim=1).view(-1, num_samples)
vllm/model_executor/layers/sampler.py
+3
-9
@@ -21,7 +21,6 @@ from vllm.sampling_params import SamplingType
21
from vllm.sequence import (VLLM_INVALID_TOKEN_ID,
22
CompletionSequenceGroupOutput, Logprob,
23
PromptLogprobs, SampleLogprobs, SequenceOutput)
24
-from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics
24
25
if envs.VLLM_USE_FLASHINFER_SAMPLER and find_spec("flashinfer"):
26
# yapf: disable
@@ -119,9 +118,6 @@ class SamplerOutput(
118
# specified in lieu of prompt token ids or text.
119
sampled_token_embeds: Optional[torch.Tensor] = None
120
122
- # Spec decode metrics populated by workers.
123
- spec_decode_worker_metrics: Optional[SpecDecodeWorkerMetrics] = None
124
-
121
# Optional last hidden states from the model.
122
hidden_states: Optional[torch.Tensor] = None
123
@@ -159,11 +155,9 @@ class SamplerOutput(
155
else self.sampled_token_probs.shape)
156
sampled_token_ids_repr = ("None" if self.sampled_token_ids is None else
157
self.sampled_token_ids.shape)
162
- return (
163
- f"SamplerOutput(outputs={self.outputs}, "
164
- f"sampled_token_probs={sampled_token_probs_repr}, "
165
- f"sampled_token_ids={sampled_token_ids_repr}, "
166
- f"spec_decode_worker_metrics={self.spec_decode_worker_metrics})")
158
+ return (f"SamplerOutput(outputs={self.outputs}, "
159
+ f"sampled_token_probs={sampled_token_probs_repr}, "
160
+ f"sampled_token_ids={sampled_token_ids_repr})")
161
162
163
class Sampler(nn.Module):
vllm/model_executor/layers/spec_decode_base_sampler.py
deleted
-259
@@ -1,259 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from abc import abstractmethod
5
-from typing import Optional, Union
6
-
7
-import torch
8
-import torch.jit
9
-import torch.nn as nn
10
-
11
-from vllm.platforms import current_platform
12
-
13
-
14
-class SpecDecodeBaseSampler(nn.Module):
15
- """Base class for samplers used for Speculative Decoding verification
16
- step.
17
- """
18
-
19
- def __init__(self, strict_mode: bool = False):
20
- """Base class constructor.
21
- Args:
22
- strict_mode: Whether or not to perform shape/device/dtype checks
23
- during sampling. This catches correctness issues but adds
24
- nontrivial latency.
25
- """
26
- super().__init__()
27
- self._strict_mode = strict_mode
28
-
29
- # NOTE: A "bonus token" is accepted iff all proposal tokens are
30
- # accepted. There is always only one possible bonus token. We store this
31
- # value in a variable for readability.
32
- self._num_bonus_tokens = 1
33
-
34
- self.num_accepted_tokens: Optional[torch.Tensor] = None
35
- self.num_emitted_tokens: Optional[torch.Tensor] = None
36
- self.num_draft_tokens: int = 0
37
-
38
- def init_gpu_tensors(self, device: Union[int, str]) -> None:
39
- assert self.num_accepted_tokens is None
40
- if isinstance(device, int):
41
- device = f"{current_platform.device_type}:{device}"
42
- elif not isinstance(device, str):
43
- raise ValueError(f"Device must be int or str, get {type(device)}")
44
- self.num_accepted_tokens = torch.tensor(0,
45
- dtype=torch.long,
46
- device=device)
47
- self.num_emitted_tokens = torch.tensor(0,
48
- dtype=torch.long,
49
- device=device)
50
-
51
- def init_tensors(self,
52
- device: Union[int, str],
53
- device_type: Union[torch.device, str] = 'cuda') -> None:
54
- assert self.num_accepted_tokens is None
55
- if isinstance(device_type, torch.device):
56
- device_type = device_type.type
57
- if isinstance(device, int):
58
- device = f"{device_type}:{device}"
59
- self.num_accepted_tokens = torch.tensor(0,
60
- dtype=torch.long,
61
- device=device)
62
- self.num_emitted_tokens = torch.tensor(0,
63
- dtype=torch.long,
64
- device=device)
65
-
66
- @property
67
- def probs_dtype(self):
68
- return torch.float32
69
-
70
- @property
71
- def token_id_dtype(self):
72
- return torch.int64
73
-
74
- def _create_output(
75
- self,
76
- accepted: torch.Tensor, # [batch_size, k]
77
- substitute_token_ids: torch.Tensor, # [batch_size, k]
78
- draft_token_ids: torch.Tensor, # [batch_size, k]
79
- bonus_token_ids: torch.Tensor, # [batch_size]
80
- ) -> torch.Tensor:
81
- """Format output. Returns a matrix of token ids. When
82
- a token is rejected via sampling, all subsequent token ids are
83
- set to -1 for the sequence.
84
-
85
- Args:
86
- accepted: A boolean tensor indicating if the corresponding
87
- draft token in draft_token_ids should be accepted or not.
88
- substitute_token_ids: A tensor of token_ids that can be used
89
- as substitutes for the draft token ids if the proposed token
90
- is rejected.
91
- draft_token_ids: A tensor of token ids speculated by the
92
- draft model.
93
- bonus_token_ids: Token ids to use as the bonus token if
94
- all the draft tokens are accepted.
95
- Returns:
96
- A tensor containing the accepted token ids. The shape of the
97
- tensor is [batch_size, k + num_bonus_tokens]
98
- """
99
- batch_size, k = substitute_token_ids.shape
100
- bonus_token_ids = bonus_token_ids.squeeze(-1)
101
- # Determine the index of the first False value for each row.
102
- limits = (accepted == 0).max(1).indices
103
- limits[~(accepted == 0).any(1)] = k
104
-
105
- # Create masks using the indices.
106
- indices = torch.arange(k, device=accepted.device).unsqueeze(0)
107
- accepted_mask = indices < limits.unsqueeze(1)
108
- after_false_mask = indices == limits.unsqueeze(1)
109
-
110
- # Create an extended output tensor
111
- output_with_bonus_tokens = -torch.ones(
112
- (batch_size, k + self._num_bonus_tokens),
113
- dtype=self.token_id_dtype,
114
- device=accepted.device)
115
- output = output_with_bonus_tokens[:, :k]
116
-
117
- # Fill in the first k columns of the output tensor using masks and data
118
- # tensors.
119
- output[:, :k] = torch.where(accepted_mask, draft_token_ids,
120
- -torch.ones_like(draft_token_ids))
121
-
122
- # Fill the last column.
123
- # We check output directly as accepted may have True values inconsistent
124
- # with causal acceptance.
125
- output_with_bonus_tokens[:, -1] = torch.where(output[:, -1] != -1,
126
- bonus_token_ids, -1)
127
-
128
- # Fill the recovered token ids.
129
- output.mul_(~after_false_mask).add_(
130
- substitute_token_ids.mul(after_false_mask))
131
-
132
- self.num_accepted_tokens += accepted.sum()
133
- self.num_emitted_tokens += (output_with_bonus_tokens != -1).sum()
134
- self.num_draft_tokens += batch_size * k
135
-
136
- return output_with_bonus_tokens
137
-
138
- def _raise_if_incorrect_input(
139
- self,
140
- target_with_bonus_probs: torch.Tensor,
141
- draft_token_ids: torch.Tensor,
142
- bonus_token_ids: torch.Tensor,
143
- draft_probs: Optional[torch.Tensor] = None,
144
- ) -> None:
145
- self._raise_if_incorrect_shape(target_with_bonus_probs,
146
- draft_token_ids, bonus_token_ids,
147
- draft_probs)
148
- self._raise_if_incorrect_dtype(target_with_bonus_probs,
149
- draft_token_ids, bonus_token_ids,
150
- draft_probs)
151
- self._raise_if_inconsistent_device(target_with_bonus_probs,
152
- draft_token_ids, bonus_token_ids,
153
- draft_probs)
154
- self._raise_if_out_of_bounds_vocab(target_with_bonus_probs.shape[-1],
155
- draft_token_ids, bonus_token_ids)
156
-
157
- def _raise_if_incorrect_shape(
158
- self,
159
- target_with_bonus_probs: torch.Tensor,
160
- draft_token_ids: torch.Tensor,
161
- bonus_token_ids: torch.Tensor,
162
- draft_probs: Optional[torch.Tensor] = None,
163
- ) -> None:
164
- (target_batch_size, num_target_probs,
165
- target_vocab_size) = target_with_bonus_probs.shape
166
-
167
- # Does not count the extra token
168
- num_target_probs -= 1
169
-
170
- # validate the shape of draft token ids.
171
- draft_token_ids_batch_size, num_draft_token_ids = draft_token_ids.shape
172
- assert draft_token_ids_batch_size == target_batch_size
173
- assert num_draft_token_ids == num_target_probs
174
-
175
- # validate the shape of bonus token ids
176
- bonus_batch_size, num_bonus_tokens = bonus_token_ids.shape
177
- assert bonus_batch_size == target_batch_size
178
- assert num_bonus_tokens == self._num_bonus_tokens
179
-
180
- # validate the shape of draft probs if it is set
181
- if draft_probs is not None:
182
- (draft_batch_size, num_draft_probs,
183
- draft_vocab_size) = draft_probs.shape
184
- assert draft_batch_size == target_batch_size
185
- assert num_draft_probs == num_target_probs
186
- assert (draft_vocab_size == target_vocab_size
187
- ), f"{draft_vocab_size=} {target_vocab_size=}"
188
-
189
- def _raise_if_incorrect_dtype(
190
- self,
191
- target_with_bonus_probs: torch.Tensor,
192
- draft_token_ids: torch.Tensor,
193
- bonus_token_ids: torch.Tensor,
194
- draft_probs: Optional[torch.Tensor] = None,
195
- ) -> None:
196
- assert target_with_bonus_probs.dtype == self.probs_dtype
197
- assert draft_token_ids.dtype == self.token_id_dtype
198
- assert bonus_token_ids.dtype == self.token_id_dtype
199
- if draft_probs is not None:
200
- assert draft_probs.dtype == self.probs_dtype
201
-
202
- def _raise_if_inconsistent_device(
203
- self,
204
- target_with_bonus_probs: torch.Tensor,
205
- draft_token_ids: torch.Tensor,
206
- bonus_token_ids: torch.Tensor,
207
- draft_probs: Optional[torch.Tensor] = None,
208
- ) -> None:
209
- devices = [
210
- t.device for t in [
211
- target_with_bonus_probs, bonus_token_ids, draft_probs,
212
- draft_token_ids
213
- ] if t is not None
214
- ]
215
- assert all([devices[0] == device for device in devices])
216
-
217
- def _raise_if_out_of_bounds_vocab(
218
- self,
219
- vocab_size: int,
220
- draft_token_ids: torch.Tensor,
221
- bonus_token_ids: torch.Tensor,
222
- ) -> None:
223
- assert torch.all(bonus_token_ids < vocab_size)
224
- assert torch.all(bonus_token_ids >= 0)
225
- assert torch.all(draft_token_ids < vocab_size)
226
- assert torch.all(draft_token_ids >= 0)
227
-
228
-
229
-class SpecDecodeDeterministicBaseSampler(SpecDecodeBaseSampler):
230
- """Base class for samplers used for Speculative Decoding verification
231
- step which are deterministic.
232
- """
233
-
234
- @abstractmethod
235
- def forward(
236
- self,
237
- target_with_bonus_probs: torch.Tensor,
238
- bonus_token_ids: torch.Tensor,
239
- draft_probs: torch.Tensor,
240
- draft_token_ids: torch.Tensor,
241
- ) -> torch.Tensor:
242
- raise NotImplementedError
243
-
244
-
245
-class SpecDecodeStochasticBaseSampler(SpecDecodeBaseSampler):
246
- """Base class for samplers used for Speculative Decoding verification
247
- step which are stochastic
248
- """
249
-
250
- @abstractmethod
251
- def forward(
252
- self,
253
- target_with_bonus_probs: torch.Tensor,
254
- bonus_token_ids: torch.Tensor,
255
- draft_probs: torch.Tensor,
256
- draft_token_ids: torch.Tensor,
257
- seeded_seqs: Optional[dict[int, torch.Generator]] = None,
258
- ) -> torch.Tensor:
259
- raise NotImplementedError
vllm/model_executor/layers/typical_acceptance_sampler.py
deleted
-166
@@ -1,166 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import torch
5
-import torch.jit
6
-
7
-from vllm.model_executor.layers.spec_decode_base_sampler import (
8
- SpecDecodeDeterministicBaseSampler)
9
-
10
-
11
-class TypicalAcceptanceSampler(SpecDecodeDeterministicBaseSampler):
12
- """Apply typical acceptance sampling as described in section 3.3.1 in
13
- "MEDUSA: Simple LLM Inference Acceleration Framework with
14
- Multiple Decoding Heads"
15
- https://arxiv.org/pdf/2401.10774
16
- """
17
-
18
- def __init__(
19
- self,
20
- posterior_threshold: float,
21
- posterior_alpha: float,
22
- strict_mode: bool = False,
23
- ):
24
- """Create a Typical Acceptance Sampler.
25
-
26
- Args:
27
- strict_mode: Whether or not to perform shape/device/dtype checks
28
- during sampling. This catches correctness issues but adds
29
- nontrivial latency.
30
- posterior_threshold : A threshold value that sets a lower bound
31
- on the posterior probability of a token in target model for it
32
- to be accepted.
33
- posterior_alpha : A scaling factor for the entropy-based
34
- threshold in typical acceptance sampling.
35
- """
36
- self._posterior_threshold = posterior_threshold
37
- self._posterior_alpha = posterior_alpha
38
- super().__init__(strict_mode=strict_mode)
39
-
40
- def forward(
41
- self,
42
- target_with_bonus_probs: torch.Tensor,
43
- bonus_token_ids: torch.Tensor,
44
- draft_probs: torch.Tensor,
45
- draft_token_ids: torch.Tensor,
46
- ) -> torch.Tensor:
47
- """Sample token ids using typical acceptance sampling. This accepts
48
- or rejects tokens proposed by the draft model using the probability
49
- of each token according to the draft and target models.
50
-
51
- In the worst case where all draft tokens are rejected, it is guaranteed
52
- one token will be emitted.
53
-
54
- In the case where all draft tokens are accepted, the bonus token will be
55
- accepted.
56
-
57
- Args:
58
- target_probs: The probability distribution over token ids given
59
- context according to the target model.
60
- shape = [batch_size, num_speculative_tokens, vocab_size]
61
-
62
- bonus_token_ids: The "bonus" token ids that are accepted iff all
63
- speculative tokens in a sequence are accepted.
64
- shape = [batch_size, num_bonus_tokens]
65
-
66
- draft_probs: This parameter is unused by the acceptance sampler.
67
-
68
- draft_token_ids: The token ids that were sampled from the draft
69
- probabilities.
70
- shape = [batch_size, num_speculative_tokens]
71
-
72
- Returns:
73
- output_token_ids: The token ids sampled via rejection sampling,
74
- or -1 if unable to sample a token because the previous token
75
- was rejected.
76
- shape = [batch_size, num_speculative_tokens + num_bonus_tokens]
77
- """
78
- # Only perform shape/dtype/device checking in strict mode, as it adds
79
- # overhead.
80
- if self._strict_mode:
81
- self._raise_if_incorrect_input(target_with_bonus_probs,
82
- draft_token_ids, bonus_token_ids)
83
- target_probs = target_with_bonus_probs[:, :-1]
84
- accepted = self._evaluate_accepted_tokens(target_probs,
85
- draft_token_ids)
86
- recovered_token_ids = self._get_recovered_token_ids(target_probs)
87
- output_token_ids = self._create_output(accepted, recovered_token_ids,
88
- draft_token_ids,
89
- bonus_token_ids)
90
- return output_token_ids
91
-
92
- def _evaluate_accepted_tokens(self, target_probs, draft_token_ids):
93
- r"""
94
- Evaluates and returns a mask of accepted tokens based on the
95
- posterior probabilities.
96
-
97
- Args:
98
- target_probs (torch.Tensor): A tensor of shape
99
- (batch_size, k, vocab_size) representing the probabilities of
100
- each token in the vocabulary for each position in the proposed
101
- sequence. This is the distribution generated by the target
102
- model.
103
- draft_token_ids (torch.Tensor): A tensor of shape (batch_size, k)
104
- representing the proposed token ids.
105
-
106
- A draft token_id x_{n+k} is accepted if it satisfies the
107
- following condition
108
-
109
- $$
110
- p_{\text{original}}(x_{n+k} | x_1, x_2, \dots, x_{n+k-1}) >
111
- \min \left( \epsilon, \delta * \exp \left(
112
- -H(p_{\text{original}}(
113
- \cdot | x_1, x_2, \ldots, x_{n+k-1})) \right) \right)
114
- $$
115
-
116
- where $p_{\text{original}}$ corresponds to target_probs
117
- and $\epsilon$ and $\delta$ correspond to hyperparameters
118
- specified using self._posterior_threshold and self._posterior_alpha
119
-
120
- This method computes the posterior probabilities for the given
121
- draft token ids based on the provided target probabilities. It
122
- calculates the entropy of the posterior distribution and determines
123
- a dynamic threshold for each token position using the provided
124
- posterior_threshold and posterior_alpha values. The method then
125
- returns a boolean mask indicating which tokens can be accepted.
126
-
127
- Returns:
128
- torch.Tensor: A boolean tensor of shape (batch_size, k) where each
129
- element indicates whether the corresponding draft token has
130
- been accepted or rejected. True indicates acceptance and false
131
- indicates rejection.
132
- """
133
- device = target_probs.device
134
- candidates_prob = torch.gather(
135
- target_probs, dim=-1,
136
- index=draft_token_ids.unsqueeze(-1)).squeeze(-1)
137
- # A small constant added to prevent computing the logarithm of zero,
138
- # which can lead to undefined values.
139
- epsilon = 1e-5
140
- posterior_entropy = -torch.sum(
141
- target_probs * torch.log(target_probs + epsilon), dim=-1)
142
- threshold = torch.minimum(
143
- torch.ones_like(posterior_entropy, device=device) *
144
- self._posterior_threshold,
145
- torch.exp(-posterior_entropy) * self._posterior_alpha,
146
- )
147
- accepted_mask = candidates_prob > threshold
148
- return accepted_mask
149
-
150
- def _get_recovered_token_ids(self, target_probs):
151
- """
152
- The recovered token ids will fill the first unmatched token
153
- by the target token.
154
-
155
- Args:
156
- target_probs (torch.Tensor): A tensor of shape
157
- (batch_size, k, vocab_size) containing the target probability
158
- distribution.
159
-
160
- Returns:
161
- torch.Tensor: A tensor of shape (batch_size, k) with the recovered
162
- token ids which are selected from target probs.
163
- """
164
- max_indices = torch.argmax(target_probs, dim=-1)
165
-
166
- return max_indices
vllm/model_executor/models/eagle.py
deleted
-261
@@ -1,261 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from collections.abc import Iterable
5
-from typing import Optional
6
-
7
-import torch
8
-import torch.nn as nn
9
-
10
-from vllm.config import VllmConfig
11
-from vllm.logger import init_logger
12
-from vllm.model_executor.layers.layernorm import RMSNorm
13
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
14
-from vllm.model_executor.layers.vocab_parallel_embedding import (
15
- DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead)
16
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
17
-from vllm.model_executor.models import ModelRegistry
18
-from vllm.model_executor.sampling_metadata import SamplingMetadata
19
-from vllm.sequence import IntermediateTensors
20
-
21
-from .utils import maybe_prefix
22
-
23
-logger = init_logger(__name__)
24
-
25
-
26
-class DummyInputLayerNorm(nn.Module):
27
-
28
- def __init__(self, weight=None, bias=None):
29
- super().__init__()
30
- self.weight = nn.Parameter(weight) if weight is not None else None
31
- self.bias = nn.Parameter(bias) if bias is not None else None
32
-
33
- def forward(self, x):
34
- return x
35
-
36
-
37
-class DummyOutputNorm(nn.Module):
38
-
39
- def forward(self, x, residual):
40
- if residual is None:
41
- return x
42
- else:
43
- return x + residual, None
44
-
45
-
46
-class EAGLE(nn.Module):
47
- """This class implements the EAGLE draft model from the paper: https://arxiv.org/pdf/2401.15077
48
- Reference implementation: https://github.com/SafeAILab/EAGLE
49
-
50
- Differences from reference implementation:
51
- 1. In reference, LlamaDecoderLayer implementation doesn't have
52
- input_layernorm for 1st decoder layer (https://github.com/SafeAILab/EAGLE/blob/7d065d084443fbfd386f88839efd7193c12be869/eagle/model/cnets.py#L427).
53
- Following this approach, our implementation also disables
54
- the input_layernorm for the first decoder layer.
55
- 2. We allow any decoder layer to be used in EAGLE whereas in reference
56
- decoder layer is fixed to be LlamaDecoderLayer.
57
- 3. We have an optional token_map which reduces draft vocab to most
58
- frequently used tokens to give some additional speed-up by reducing
59
- sampling overhead. This is disabled unless the checkpoint file has
60
- explicit token_map tensor and config has an optional attribute
61
- truncated_vocab_size < vocab_size. To use this technique, one has to find
62
- the top-k most frequent tokens in target dataset and add that as a tensor
63
- in the draft checkpoint (using key token_map). Also, the draft config
64
- needs to have truncated_vocab_size (=k) as an attribute.
65
- 4. We allow an enhanced EAGLE architecture similar to the DeepSeek MTP
66
- module with regards to the use of additional RMS norms. The original
67
- EAGLE architecture 1) skips the pre-attention norm in its first
68
- transformer block, and 2) skips the final output norm, both of which we
69
- found to be suboptimal. We also add the support for separate norms
70
- applying to both the token embedding and hidden states before projection
71
- as in DeepSeek MTP, which we found to improve performance as well.
72
- """
73
-
74
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
75
- super().__init__()
76
- config = vllm_config.model_config.hf_config
77
- self.dtype = vllm_config.model_config.dtype
78
- self.config = config
79
-
80
- architectures = getattr(self.config.model, "architectures", [])
81
- model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
82
-
83
- self.model = model_cls(vllm_config=vllm_config,
84
- prefix=maybe_prefix(prefix, "model"))
85
-
86
- self.fc = nn.Linear(config.model.hidden_size * 2,
87
- config.model.hidden_size,
88
- bias=getattr(self.config, "eagle_fc_bias", False))
89
-
90
- # Modify layer normalization and residual connections as suggested
91
- # in the EAGLE framework: https://github.com/SafeAILab/EAGLE
92
- # While weights and biases are generally not needed,
93
- # they are retained here to support certain unit tests
94
- # (e.g., spec_decode/e2e/test_eagle_correctness.py).
95
- if not hasattr(self.config.model,
96
- "skip_prenorm") or self.config.model.skip_prenorm:
97
- self.model.model.layers[0].input_layernorm = DummyInputLayerNorm(
98
- weight=self.model.model.layers[0].input_layernorm.weight)
99
-
100
- if not hasattr(
101
- self.config.model,
102
- "skip_output_norm") or self.config.model.skip_output_norm:
103
- self.model.model.norm = DummyOutputNorm()
104
-
105
- self.add_para_norm = False
106
- if hasattr(self.config.model,
107
- "add_para_norm") and self.config.model.add_para_norm:
108
- self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
109
- self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
110
- self.add_para_norm = True
111
-
112
- self.orig_vocab_size = config.vocab_size
113
- self.truncated_vocab_size = config.truncated_vocab_size
114
- self.unpadded_vocab_size = self.truncated_vocab_size
115
-
116
- self.lm_head = ParallelLMHead(
117
- self.unpadded_vocab_size,
118
- config.hidden_size,
119
- org_num_embeddings=self.truncated_vocab_size,
120
- padding_size=DEFAULT_VOCAB_PADDING_SIZE,
121
- )
122
-
123
- logit_scale = getattr(config, "logit_scale", 1.0)
124
- self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
125
- self.truncated_vocab_size,
126
- logit_scale)
127
-
128
- # Token map is a idx to token mapping to reduce the vocab size for
129
- # the draft model. Using smaller vocab size for draft, containing
130
- # only most frequent tokens reduces the speculation overhead. This
131
- # doesn't affect the acceptance rate much and thus gives more speed
132
- # -up. By default, this is disabled and is only used if the EAGLE
133
- # checkpoint file has token_map tensor.
134
- self.token_map = None
135
-
136
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
137
- return self.model.model.get_input_embeddings(input_ids)
138
-
139
- def forward(
140
- self,
141
- input_ids: torch.Tensor,
142
- positions: torch.Tensor,
143
- previous_hidden_states: torch.Tensor,
144
- intermediate_tensors: Optional[IntermediateTensors] = None,
145
- inputs_embeds: Optional[torch.Tensor] = None,
146
- ) -> torch.Tensor:
147
-
148
- if inputs_embeds is None:
149
- inputs_embeds = self.get_input_embeddings(input_ids)
150
-
151
- # Handle both empty previous_hidden_states
152
- # and mismatched batch size
153
- batch_size = inputs_embeds.size(0)
154
- if previous_hidden_states.size(0) == 0 or \
155
- previous_hidden_states.size(0) != batch_size:
156
- hidden_dim = self.config.model.hidden_size
157
- device = inputs_embeds.device
158
- # Create zero tensor with matching batch size
159
- previous_hidden_states = \
160
- torch.zeros(batch_size, hidden_dim, device=device)
161
-
162
- if self.add_para_norm:
163
- inputs_embeds = torch.cat([
164
- self.enorm(inputs_embeds),
165
- self.hnorm(previous_hidden_states)
166
- ],
167
- dim=-1)
168
- else:
169
- inputs_embeds = torch.cat([inputs_embeds, previous_hidden_states],
170
- dim=-1)
171
-
172
- inputs_embeds = self.fc(inputs_embeds)
173
-
174
- inputs_embeds[positions == 0] = 0 # masking inputs at position=0
175
-
176
- hidden_states = self.model.model(
177
- input_ids=None,
178
- inputs_embeds=inputs_embeds,
179
- positions=positions,
180
- intermediate_tensors=intermediate_tensors,
181
- )
182
- return hidden_states
183
-
184
- def compute_logits(self, hidden_states: torch.Tensor,
185
- sampling_metadata: SamplingMetadata) -> torch.Tensor:
186
- logits = self.logits_processor(self.lm_head, hidden_states,
187
- sampling_metadata)
188
-
189
- if self.token_map is not None:
190
- _logits = logits
191
- logits = -torch.inf * torch.ones(
192
- size=(*_logits.shape[:-1], self.orig_vocab_size),
193
- device=_logits.device,
194
- dtype=_logits.dtype)
195
-
196
- logits[..., self.token_map] = _logits
197
-
198
- return logits
199
-
200
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
201
- # This implementation is incompatible with https://huggingface.co/yuhuili/EAGLE-LLaMA3-Instruct-8B
202
- # due to missing lm_head weights and its config being that of a
203
- # Llama model. Here's a compatible version with the same weights:
204
- # https://huggingface.co/abhigoyal/EAGLE-LLaMA3-Instruct-8B-vllm
205
- # Also, here's an example script for converting trained EAGLE
206
- # checkpoint to vLLM compatible version: https://gist.github.com/abhigoyal1997/1e7a4109ccb7704fbc67f625e86b2d6d
207
- model_weights = {}
208
- for name, loaded_weight in weights:
209
- if name == "token_map":
210
- if self.config.truncated_vocab_size < self.config.vocab_size:
211
- self.token_map = nn.Parameter(loaded_weight,
212
- requires_grad=False)
213
- elif name.startswith("fc.weight"):
214
- weight_loader = getattr(self.fc.weight, "weight_loader",
215
- default_weight_loader)
216
- weight_loader(self.fc.weight, loaded_weight)
217
- elif name.startswith("fc.bias"):
218
- if self.fc.bias is not None:
219
- weight_loader = getattr(self.fc.bias, "weight_loader",
220
- default_weight_loader)
221
- weight_loader(self.fc.bias, loaded_weight)
222
- else:
223
- logger.warning_once("Found bias in the loaded weights but "
224
- "the model config doesn't have bias.")
225
- elif name.startswith("enorm.weight"):
226
- weight_loader = getattr(self.enorm.weight, "weight_loader",
227
- default_weight_loader)
228
- weight_loader(self.enorm.weight, loaded_weight)
229
- elif name.startswith("hnorm.weight"):
230
- weight_loader = getattr(self.hnorm.weight, "weight_loader",
231
- default_weight_loader)
232
- weight_loader(self.hnorm.weight, loaded_weight)
233
- elif name.startswith("model.lm_head.") or name.startswith(
234
- "model.model."):
235
- model_weights[name.split("model.", 1)[-1]] = loaded_weight
236
- elif name.startswith("lm_head.") or name.startswith("model."):
237
- model_weights[name] = loaded_weight
238
- else:
239
- model_weights[f"model.{name}"] = loaded_weight
240
-
241
- if "lm_head.weight" in model_weights:
242
- lm_head_weight = model_weights.pop("lm_head.weight")
243
-
244
- if self.token_map is not None and\
245
- lm_head_weight.shape[0] > self.token_map.shape[0]:
246
-
247
- lm_head_weight = lm_head_weight[self.token_map]
248
-
249
- else:
250
- # NOTE(Shangming): initialize the placeholder for lm_head weight.
251
- lm_head_weight = torch.zeros(
252
- self.lm_head.org_vocab_size,
253
- self.lm_head.embedding_dim,
254
- dtype=self.dtype,
255
- )
256
-
257
- weight_loader = getattr(self.lm_head.weight, "weight_loader",
258
- default_weight_loader)
259
- weight_loader(self.lm_head.weight, lm_head_weight)
260
-
261
- self.model.load_weights(model_weights.items())
vllm/model_executor/models/registry.py
+3
-2
@@ -239,14 +239,15 @@ _MULTIMODAL_MODELS = {
239
240
_SPECULATIVE_DECODING_MODELS = {
241
"MiMoMTPModel": ("mimo_mtp", "MiMoMTP"),
242
- "EAGLEModel": ("eagle", "EAGLE"),
242
"EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),
243
"EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
244
"EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
245
"Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
246
"DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
247
"MedusaModel": ("medusa", "Medusa"),
249
- "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
248
+ # Temporarily disabled.
249
+ # # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.
250
+ # "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
251
}
252
253
_TRANSFORMERS_MODELS = {
vllm/platforms/cuda.py
+4
-8
@@ -132,14 +132,10 @@ class CudaPlatformBase(Platform):
132
parallel_config.worker_cls = \
133
"vllm.worker.multi_step_worker.MultiStepWorker"
134
elif vllm_config.speculative_config:
135
- if envs.VLLM_USE_V1:
136
- parallel_config.worker_cls = \
137
- "vllm.v1.worker.gpu_worker.Worker"
138
- else:
139
- parallel_config.worker_cls = \
140
- "vllm.spec_decode.spec_decode_worker.create_spec_worker"
141
- parallel_config.sd_worker_cls = \
142
- "vllm.worker.worker.Worker"
135
+ if not envs.VLLM_USE_V1:
136
+ raise NotImplementedError(
137
+ "Speculative decoding is not supported on vLLM V0.")
138
+ parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
139
else:
140
if envs.VLLM_USE_V1:
141
parallel_config.worker_cls = \
vllm/platforms/rocm.py
+3
-8
@@ -326,15 +326,10 @@ class RocmPlatform(Platform):
326
parallel_config.worker_cls = \
327
"vllm.worker.multi_step_worker.MultiStepWorker"
328
elif vllm_config.speculative_config:
329
- if envs.VLLM_USE_V1:
329
+ if not envs.VLLM_USE_V1:
330
raise NotImplementedError(
331
- "Speculative decoding is not yet supported on vLLM V1."
332
- )
333
- else:
334
- parallel_config.worker_cls = \
335
- "vllm.spec_decode.spec_decode_worker.create_spec_worker"
336
- parallel_config.sd_worker_cls = \
337
- "vllm.worker.worker.Worker"
331
+ "Speculative decoding is not supported on vLLM V0.")
332
+ parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
333
else:
334
if envs.VLLM_USE_V1:
335
parallel_config.worker_cls = \
vllm/sequence.py
+1
-13
@@ -112,13 +112,6 @@ class RequestMetrics:
112
model_execute_time: The time spent in the model execute function. This
113
will include model forward, block/sync across
114
workers, cpu-gpu sync time and sampling time.
115
- spec_token_acceptance_counts: number of accepted speculative tokens at
116
- each position; the first token is from
117
- the target model and is always accepted;
118
- e.g., when it's [10, 8, 4, 2] for a req,
119
- it means there were 10 forward passes in
120
- total, and there were 8, 4, 2 accepted
121
- tokens at 1st, 2nd, 3rd speculation step.
115
"""
116
arrival_time: float
117
last_token_time: float
@@ -129,7 +122,6 @@ class RequestMetrics:
122
scheduler_time: Optional[float] = None
123
model_forward_time: Optional[float] = None
124
model_execute_time: Optional[float] = None
132
- spec_token_acceptance_counts: Optional[list[int]] = None
125
126
127
class SequenceDataDelta(
@@ -748,9 +740,7 @@ class SequenceGroup:
740
last_token_time=arrival_time,
741
first_scheduled_time=None,
742
first_token_time=None,
751
- time_in_queue=None,
752
- spec_token_acceptance_counts=[0] *
753
- draft_size)
743
+ time_in_queue=None)
744
self.last_token_latency = 0.0
745
self.lora_request = lora_request
746
self.prompt_logprobs: Optional[PromptLogprobs] = None
@@ -1390,8 +1380,6 @@ class ExecuteModelRequest(
1380
previous_hidden_states: Optional[HiddenStates] = None
1381
# The number of forward steps to run.
1382
num_steps: int = 1
1393
- # The step index for spec model input.
1394
- spec_step_idx: Optional[int] = None
1383
# Finished request ids since last step.
1384
finished_requests_ids: list[str] = msgspec.field(default_factory=list)
1385
# The last sampled token ids for multi step decoding.
vllm/spec_decode/__init__.py
vllm/spec_decode/batch_expansion.py
deleted
-506
@@ -1,506 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from array import array
5
-from itertools import chain, count
6
-from typing import Iterator, List, Optional, Tuple
7
-
8
-import torch
9
-
10
-from vllm import SamplingParams
11
-from vllm.model_executor.layers.sampler import SamplerOutput
12
-from vllm.sequence import (VLLM_INVALID_TOKEN_ID, VLLM_TOKEN_ID_ARRAY_TYPE,
13
- ExecuteModelRequest, SequenceData,
14
- SequenceGroupMetadata, get_all_seq_ids)
15
-from vllm.spec_decode.interfaces import (SpeculativeProposals,
16
- SpeculativeScorer, SpeculativeScores)
17
-from vllm.spec_decode.util import nvtx_range, split_batch_by_proposal_len
18
-
19
-SeqId = int
20
-TargetSeqId = int
21
-TokenId = int
22
-
23
-DEFAULT_SIMPLE_SAMPLING_PARAMS = SamplingParams()
24
-
25
-
26
-class BatchExpansionTop1Scorer(SpeculativeScorer):
27
- """Implements a speculative scorer that uses batch expansion to get
28
- probabilities of speculative tokens according to the scoring model.
29
-
30
- Batch expansion converts a list of sequences and multiple query positions
31
- to a new batch of sequences, each with a single query position. This allows
32
- for MQA-like scoring in speculative decoding without requiring an MQA
33
- kernel.
34
-
35
- It is strictly less efficient than MQA scoring.
36
-
37
- It only supports scoring the top1 proposal tokens of the proposer, instead
38
- of topk/tree.
39
- """
40
-
41
- @nvtx_range("BatchExpansionTop1Scorer.score_proposals")
42
- def score_proposals(
43
- self,
44
- execute_model_req: ExecuteModelRequest,
45
- proposals: SpeculativeProposals,
46
- ) -> SpeculativeScores:
47
- """Score the proposed tokens via the scorer model.
48
-
49
- This converts each input sequence to a set of k+1 target sequences. The
50
- target sequences have the unique continuations to be scored and a
51
- unique sequence ID that is different from all input sequence ids.
52
-
53
- If a speculative sequence length would exceed the max model length, then
54
- no speculation is produced for that sequence.
55
-
56
- Args:
57
- execute_model_req: The execution request.
58
- proposals: The speculative proposals to score.
59
- Returns:
60
- SpeculativeScores: The scores of each speculative token, along with
61
- which sequences were ignored during scoring.
62
- """
63
-
64
- # TODO(cade) perform this on GPU to remove blocking call.
65
- proposal_lens_list = proposals.proposal_lens.tolist()
66
- proposal_token_ids_list = proposals.proposal_token_ids.tolist()
67
-
68
- # Filter the list to ignore invalid proposals.
69
- proposal_token_ids_list_without_skips = [
70
- proposals for proposals in proposal_token_ids_list
71
- if VLLM_INVALID_TOKEN_ID not in proposals
72
- ]
73
-
74
- (spec_indices, non_spec_indices, target_seq_group_metadata_list,
75
- num_scoring_tokens) = self._expand_batch(
76
- seq_group_metadata_list=execute_model_req.seq_group_metadata_list,
77
- proposal_token_ids_list=proposal_token_ids_list_without_skips,
78
- proposal_lens_list=proposal_lens_list,
79
- )
80
-
81
- target_sampler_output = self._scorer_worker.execute_model(
82
- execute_model_req=execute_model_req.clone(
83
- seq_group_metadata_list=target_seq_group_metadata_list))
84
- assert len(target_sampler_output) == 1, "expected single-step output"
85
- target_sampler_output = target_sampler_output[0]
86
-
87
- if not non_spec_indices:
88
- # All sequence groups in batch have spec decoding enabled
89
- return self._contract_batch_all_spec(
90
- target_sampler_output=target_sampler_output,
91
- proposals=proposals,
92
- )
93
- else:
94
- # Batch has a mix of spec decode enabled and disabled seq groups
95
- return self._contract_batch(
96
- execute_model_req.seq_group_metadata_list,
97
- target_sampler_output=target_sampler_output,
98
- proposals=proposals,
99
- num_scoring_tokens=num_scoring_tokens,
100
- non_spec_indices=non_spec_indices,
101
- spec_indices=spec_indices,
102
- k=execute_model_req.num_lookahead_slots,
103
- )
104
-
105
- def _expand_batch(
106
- self,
107
- seq_group_metadata_list: List[SequenceGroupMetadata],
108
- proposal_token_ids_list: List[List[TokenId]],
109
- proposal_lens_list: List[int],
110
- ) -> Tuple[List[int], List[int], List[SequenceGroupMetadata], int]:
111
- """Given the input sequences and potentially multiple corresponding
112
- proposal tokens, create a new batch where each sequence has a single
113
- query token.
114
- """
115
-
116
- # vLLM currently only supports proposal lens equal to zero or the batch
117
- # proposal len. This adds some complexity (splitting the batch into spec
118
- # and non spec sequences) and should be removed in the future. It can be
119
- # done by supporting per-sequence proposal lens.
120
- (spec_seqs, spec_indices), (non_spec_seqs, non_spec_indices) = \
121
- split_batch_by_proposal_len(
122
- seq_group_metadata_list, proposal_lens_list)
123
-
124
- spec_expanded_seqs = self._create_scoring_model_input(
125
- seq_group_metadata_list=spec_seqs,
126
- proposal_token_ids=proposal_token_ids_list,
127
- # NOTE: We determine the seq ids in the expanded batch using the
128
- # full seq_group_metadata_list, instead of only spec_seqs.
129
- target_seq_ids_iter=self._create_target_seq_id_iterator(
130
- seq_ids=get_all_seq_ids(seq_group_metadata_list)),
131
- )
132
-
133
- num_scoring_tokens = len(spec_expanded_seqs)
134
- # Batch speculative and non-speculative (e.g. chunked prefill) requests
135
- # but make sure order is prefill|decode due to backend requirement.
136
- target_seq_group_metadata_list = non_spec_seqs + spec_expanded_seqs
137
-
138
- return (spec_indices, non_spec_indices, target_seq_group_metadata_list,
139
- num_scoring_tokens)
140
-
141
- def _contract_non_speculative(
142
- self, scores: SpeculativeScores,
143
- seq_group_metadata_list: List[SequenceGroupMetadata],
144
- non_spec_indices: List[int], non_spec_outputs: SpeculativeScores,
145
- has_prompt_log: bool) -> SpeculativeScores:
146
- """
147
- Augment input `scores` with non-speculative requests outputs.
148
- This includes decode requests with speculation turned off, as well
149
- as prefill requests when `enable_chunked_prefill` is set.
150
- For the latter, prefills are further separated into terminal and
151
- non-terminal chunks (from which no token is sampled).
152
- """
153
- if not non_spec_indices:
154
- return scores
155
-
156
- if has_prompt_log:
157
- # When prompt_logprobs is enabled, prefills yield output token
158
- # (and respective prob) in the last entry (prompt|out):
159
- # [.|.|.|prefill0_out|.|prefill1_out|decode0_out|..].
160
- # With chunked prefill, non-terminal chunks have -1 on each
161
- # position: they're still picked, but they're discarded later.
162
- seq_meta = seq_group_metadata_list
163
- nospec_sizes = torch.tensor([
164
- seq_meta[i].token_chunk_size if seq_meta[i].is_prompt else 1
165
- for i in non_spec_indices
166
- ])
167
- nospec_sampled_token_idxs = torch.cumsum(nospec_sizes, 0).add_(-1)
168
- else:
169
- # In this case only sampled tokens are returned, select all.
170
- nospec_sampled_token_idxs = list(
171
- range(len(non_spec_outputs.token_ids)))
172
-
173
- scores.token_ids[non_spec_indices, :1] = \
174
- non_spec_outputs.token_ids[nospec_sampled_token_idxs].unsqueeze(1)
175
- scores.probs[non_spec_indices, :1, :] = \
176
- non_spec_outputs.probs[nospec_sampled_token_idxs].unsqueeze(1)
177
- scores.logprobs[non_spec_indices, :1, :] = \
178
- non_spec_outputs.logprobs[nospec_sampled_token_idxs].unsqueeze(1)
179
- if scores.hidden_states is not None:
180
- assert non_spec_outputs.hidden_states is not None
181
- scores.hidden_states[non_spec_indices, :1, :] = \
182
- non_spec_outputs.hidden_states[nospec_sampled_token_idxs].unsqueeze(1)
183
- return scores
184
-
185
- def _contract_batch(
186
- self,
187
- contracted_seq_group_metadata_list: List[SequenceGroupMetadata],
188
- target_sampler_output: SamplerOutput,
189
- proposals: SpeculativeProposals, num_scoring_tokens: int,
190
- non_spec_indices: List[int], spec_indices: List[int],
191
- k: int) -> SpeculativeScores:
192
- """Contract the expanded batch back into its original size.
193
- This maps the scores of speculative tokens back to their original
194
- sequences.
195
-
196
- contracted_bs is the original batch size, and the batch size that the
197
- target_sampler_output will be contracted to.
198
- """
199
- contracted_bs = len(contracted_seq_group_metadata_list)
200
- (target_token_ids, target_probs, target_logprobs, target_hidden_states,
201
- non_spec_target_token_ids, non_spec_target_probs,
202
- non_spec_target_logprobs,
203
- non_spec_target_hidden_states) = self._split_scoring_output(
204
- target_sampler_output, num_scoring_tokens)
205
-
206
- # Map distinct sequences used to score each token
207
- # of shape [batch_size * k + 1] back to [batch_size, k + 1].
208
- expanded_batch_size, k = proposals.proposal_token_ids.shape
209
-
210
- # The number of tokens in the expanded batch used for speculation is
211
- # equal to the total expanded batch size minus the number of samples for
212
- # non-speculative sequences, prefill chunks with no out tokens included
213
- non_spec_expanded_bs = len(non_spec_indices)
214
- spec_expanded_bs = expanded_batch_size - non_spec_expanded_bs
215
-
216
- target_token_ids = target_token_ids.reshape(spec_expanded_bs, k + 1)
217
- target_probs = target_probs.reshape(*target_token_ids.shape,
218
- self._vocab_size)
219
- target_logprobs = target_logprobs.reshape(target_probs.shape)
220
-
221
- if target_hidden_states is not None:
222
- target_hidden_states = target_hidden_states.reshape(
223
- *target_token_ids.shape, target_hidden_states.shape[-1])
224
-
225
- all_tokens = target_token_ids.new_full(size=(contracted_bs, k + 1),
226
- fill_value=-1)
227
- all_probs = target_probs.new_zeros(*all_tokens.shape, self._vocab_size)
228
- all_logprobs = target_logprobs.new_full(size=all_probs.shape,
229
- fill_value=-float("inf"))
230
-
231
- if target_sampler_output.hidden_states is not None:
232
- all_hidden_states = target_hidden_states.new_zeros(
233
- size=(contracted_bs, k + 1, target_hidden_states.shape[-1]))
234
- else:
235
- all_hidden_states = None
236
-
237
- has_prompt_log = any((sg.sampling_params.prompt_logprobs
238
- and sg.sampling_params.prompt_logprobs > 0)
239
- for sg in contracted_seq_group_metadata_list)
240
- # When prompt logprobs is enabled, lens of returned tensors go from
241
- # n_sampled (requests with do_sample=True) to n_prompt+n_prefills.
242
- # We adjust stride accordingly to get the generated tokens and
243
- # their probs, but pass on prompt_logprobs as is.
244
- prompt_logprobs = None
245
- if (not self._scorer_worker.model_runner.disable_logprobs\
246
- and has_prompt_log):
247
- prompt_logprobs = [
248
- o.prompt_logprobs for o in target_sampler_output.outputs
249
- ]
250
- elif not has_prompt_log:
251
- # When prompt logprobs are not to be returned,
252
- # we can ignore non-terminal chunks (no out token).
253
- non_spec_indices = [
254
- idx for idx in non_spec_indices
255
- if contracted_seq_group_metadata_list[idx].do_sample
256
- ]
257
-
258
- # "Contract" speculative.
259
- if spec_indices:
260
- all_tokens[spec_indices] = target_token_ids
261
- all_probs[spec_indices] = target_probs
262
- all_logprobs[spec_indices] = target_logprobs
263
- if all_hidden_states is not None:
264
- all_hidden_states[spec_indices] = target_hidden_states
265
-
266
- spec_scores = SpeculativeScores(probs=all_probs,
267
- token_ids=all_tokens,
268
- logprobs=all_logprobs,
269
- hidden_states=all_hidden_states,
270
- prompt_logprobs=prompt_logprobs)
271
-
272
- non_spec_outputs = SpeculativeScores(
273
- probs=non_spec_target_probs,
274
- token_ids=non_spec_target_token_ids,
275
- logprobs=non_spec_target_logprobs,
276
- hidden_states=non_spec_target_hidden_states)
277
- # Contract remaining nonspec entries based on non_spec_indices, if any.
278
- return self._contract_non_speculative(
279
- spec_scores, contracted_seq_group_metadata_list, non_spec_indices,
280
- non_spec_outputs, has_prompt_log)
281
-
282
- def _contract_batch_all_spec(
283
- self,
284
- target_sampler_output: SamplerOutput,
285
- proposals: SpeculativeProposals,
286
- ) -> SpeculativeScores:
287
- """Contract the expanded batch back into its original size.
288
- This maps the scores of speculative tokens back to their original
289
- sequences.
290
-
291
- It assumes all sequences in the batch were previously expanded.
292
- """
293
-
294
- # Map distinct sequences used to score each token
295
- # of shape [batch_size * k + 1] back to [batch_size, k + 1].
296
- contracted_bs, k = proposals.proposal_token_ids.shape
297
-
298
- # Reshape tensors to original batch size
299
- target_token_ids = target_sampler_output.sampled_token_ids.reshape(
300
- contracted_bs, k + 1)
301
- target_probs = target_sampler_output.sampled_token_probs.reshape(
302
- *target_token_ids.shape, self._vocab_size)
303
- target_logprobs = target_sampler_output.logprobs.reshape(
304
- target_probs.shape)
305
- target_hidden_states = target_sampler_output.hidden_states
306
- if target_hidden_states is not None:
307
- target_hidden_states = target_hidden_states.reshape(
308
- *target_token_ids.shape, target_hidden_states.shape[-1])
309
-
310
- return SpeculativeScores(probs=target_probs,
311
- token_ids=target_token_ids,
312
- logprobs=target_logprobs,
313
- hidden_states=target_hidden_states,
314
- prompt_logprobs=None)
315
-
316
- def _create_scoring_model_input(
317
- self,
318
- seq_group_metadata_list: List[SequenceGroupMetadata],
319
- proposal_token_ids: List[List[TokenId]], # shape: [batch_size, k]
320
- target_seq_ids_iter: Iterator[TargetSeqId],
321
- ) -> List[SequenceGroupMetadata]:
322
- """Given the original input sequences and proposed tokens from the draft
323
- model, create a list of target sequences that can be used for scoring.
324
-
325
- target_seq_ids_iter provides sequence ids for the expanded batch,
326
- fulfilling the requirement that no seq id in the expanded batch is equal
327
- to the seq id in the original batch.
328
- """
329
-
330
- if not seq_group_metadata_list:
331
- return []
332
-
333
- target_seq_group_metadata = list(
334
- chain.from_iterable(
335
- self._create_target_seq_group_metadata(
336
- seq_group_metadata,
337
- proposal_token_ids,
338
- i,
339
- target_seq_ids_iter,
340
- ) for i, seq_group_metadata in enumerate(
341
- seq_group_metadata_list)))
342
-
343
- return target_seq_group_metadata
344
-
345
- def _create_target_seq_group_metadata(
346
- self,
347
- input_seq_group_metadata: SequenceGroupMetadata,
348
- proposal_token_ids: List[List[TokenId]], # shape: [batch_size, k]
349
- batch_index: int,
350
- target_seq_ids_iter: Iterator[TargetSeqId],
351
- ) -> List[SequenceGroupMetadata]:
352
- """Given an input sequence group metadata and a list of draft tokens,
353
- create a list of target SequenceGroupMetadata, one for each
354
- token id that needs to be scored.
355
-
356
- Naive speculative decoding requires K target model scores, one for each
357
- draft model token. However one can add a bonus token such that if each
358
- token is accepted, then a final token may be sampled from the model.
359
- This function creates K+1 target SequenceGroupMetadata to take
360
- advantage of the bonus token.
361
- """
362
- assert len(input_seq_group_metadata.seq_data) == 1, (
363
- "Beam search "
364
- "not supported in speculative decoding")
365
- input_seq_id = next(iter(input_seq_group_metadata.seq_data.keys()))
366
-
367
- token_ids_to_score = self._get_token_ids_to_score(
368
- proposal_token_ids[batch_index])
369
-
370
- sampling_params = input_seq_group_metadata.sampling_params
371
- target_seq_group_metadata_list: List[SequenceGroupMetadata] = []
372
- for i, token_ids in enumerate(token_ids_to_score):
373
- target_seq_group_metadata_list.append(
374
- self._create_single_target_seq_group_metadata(
375
- input_seq_group_metadata,
376
- input_seq_id,
377
- next(target_seq_ids_iter),
378
- token_ids,
379
- sampling_params=sampling_params,
380
- ))
381
-
382
- return target_seq_group_metadata_list
383
-
384
- @staticmethod
385
- def _create_single_target_seq_group_metadata(
386
- seq_group_metadata: SequenceGroupMetadata,
387
- seq_id: SeqId,
388
- target_seq_id: TargetSeqId,
389
- token_ids: List[TokenId],
390
- sampling_params: SamplingParams,
391
- ) -> SequenceGroupMetadata:
392
- """Create a single target SequenceGroupMetadata.
393
-
394
- Args:
395
- seq_group_metadata: The metadata for the input sequence.
396
- seq_id: The input sequence ID.
397
- target_seq_id: The corresponding target sequence ID.
398
- token_ids: The list of token ids that are to be appended to the
399
- input sequence.
400
- """
401
- seq_data = seq_group_metadata.seq_data[seq_id]
402
- prompt_token_ids = seq_data.prompt_token_ids_array
403
- new_output_token_ids = [*seq_data.get_output_token_ids(), *token_ids]
404
- mrope_position_delta = seq_data.mrope_position_delta
405
-
406
- new_seq_data_dict = {
407
- target_seq_id:
408
- SequenceData(
409
- prompt_token_ids,
410
- _output_token_ids=array(VLLM_TOKEN_ID_ARRAY_TYPE,
411
- new_output_token_ids),
412
- ),
413
- }
414
- # This is a hack. Technically, spec decoding should compute
415
- # num_lookahead slots at one shot, but instead, it expands the batch
416
- # and evaluate one by one right now. context_len is seq_len - 1 because
417
- # the kv cache is filled by a previous batch in the batch expansion.
418
- for data in new_seq_data_dict.values():
419
- data.update_num_computed_tokens(data.get_len() - 1)
420
- data.mrope_position_delta = mrope_position_delta
421
-
422
- return SequenceGroupMetadata(
423
- request_id=seq_group_metadata.request_id,
424
- is_prompt=seq_group_metadata.is_prompt,
425
- seq_data=new_seq_data_dict,
426
- sampling_params=sampling_params,
427
- block_tables={
428
- target_seq_id: seq_group_metadata.block_tables[seq_id],
429
- },
430
- lora_request=None,
431
- token_chunk_size=1,
432
- )
433
-
434
- @staticmethod
435
- def _split_scoring_output(
436
- sampler_output: SamplerOutput, num_scoring_tokens: int
437
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor,
438
- Optional[torch.Tensor], torch.Tensor, torch.Tensor,
439
- torch.Tensor, Optional[torch.Tensor]]:
440
- """Split the target model output into speculative and non-speculative
441
- output.
442
- """
443
-
444
- # vLLM currently only supports proposal lens equal to zero or the batch
445
- # proposal len. This adds some complexity (splitting the batch into spec
446
- # and non spec sequences) and should be removed in the future. It can be
447
- # done by supporting per-sequence proposal lens.
448
- #
449
- # First samples are non-speculative, latter samples are from speculative
450
- # scoring (prefill|decode order).
451
- split_sizes = (sampler_output.sampled_token_ids.numel() -
452
- num_scoring_tokens, num_scoring_tokens)
453
- (non_spec_probs,
454
- spec_probs) = sampler_output.sampled_token_probs.split(split_sizes)
455
- (non_spec_sampled_tokens, spec_sampled_tokens
456
- ) = sampler_output.sampled_token_ids.flatten().split(split_sizes)
457
- (non_spec_logprobs,
458
- spec_logprobs) = sampler_output.logprobs.split(split_sizes)
459
-
460
- if sampler_output.hidden_states is not None:
461
- (non_spec_hidden_states, spec_hidden_states
462
- ) = sampler_output.hidden_states.split(split_sizes)
463
- else:
464
- non_spec_hidden_states, spec_hidden_states = None, None
465
-
466
- return (spec_sampled_tokens, spec_probs, spec_logprobs,
467
- spec_hidden_states, non_spec_sampled_tokens, non_spec_probs,
468
- non_spec_logprobs, non_spec_hidden_states)
469
-
470
- @staticmethod
471
- def _create_target_seq_id_iterator(
472
- seq_ids: List[SeqId]) -> Iterator[TargetSeqId]:
473
- """Create an iterator for creating target sequence ids.
474
- Target sequence ids are distinct from sequence ids because we create a
475
- distinct target sequence id for each proposal token to be scored.
476
-
477
- This implementation increments a counter starting at 1 + max of all
478
- provided input sequence ids.
479
- """
480
- return count(start=max(seq_ids) + 1)
481
-
482
- @staticmethod
483
- def _get_token_ids_to_score(
484
- full_spec_token_ids: List[TokenId] # shape: [k]
485
- ) -> List[List[TokenId]]:
486
- """Given an int tensor of proposal token ids, return a list of
487
- token ids that should be scored.
488
-
489
- Returns k+1 output lists. The additional one is used for generating the
490
- bonus token.
491
-
492
- Example:
493
- Input: [0, 1, 2, 3] (k=4)
494
- Output: (k+1 lists)
495
- []
496
- [0]
497
- [0, 1]
498
- [0, 1, 2]
499
- [0, 1, 2, 3]
500
- """
501
- empty_token_ids: List[TokenId] = []
502
-
503
- token_ids_to_score = [empty_token_ids]
504
- token_ids_to_score.extend(full_spec_token_ids[:i + 1]
505
- for i in range(len(full_spec_token_ids)))
506
- return token_ids_to_score
vllm/spec_decode/draft_model_runner.py
deleted
-349
@@ -1,349 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from typing import List, Optional
5
-
6
-import torch
7
-
8
-from vllm.forward_context import set_forward_context
9
-from vllm.model_executor.layers.sampler import SamplerOutput
10
-
11
-try:
12
- try:
13
- from vllm.attention.backends.flash_attn import FlashAttentionMetadata
14
- except (ModuleNotFoundError, ImportError):
15
- # vllm_flash_attn is not installed, try the ROCm FA metadata
16
- from vllm.attention.backends.rocm_flash_attn import (
17
- ROCmFlashAttentionMetadata as FlashAttentionMetadata)
18
-except (ModuleNotFoundError, ImportError) as err:
19
- raise RuntimeError(
20
- "Draft model speculative decoding currently only supports "
21
- "CUDA and ROCm flash attention backend.") from err
22
-
23
-from vllm.logger import init_logger
24
-from vllm.multimodal import MultiModalKwargs
25
-from vllm.sequence import ExecuteModelRequest, IntermediateTensors
26
-from vllm.worker.model_runner_base import (ModelRunnerBase,
27
- ModelRunnerInputBase,
28
- ModelRunnerWrapperBase)
29
-
30
-logger = init_logger(__name__)
31
-
32
-# A flag to enable debug prints for the updated input tensors
33
-# before each step.
34
-debug_advance_input = False
35
-# A flag to allow GPU advance step for draft model runner.
36
-# Set to False for debugging.
37
-allow_gpu_advance_step = True
38
-
39
-
40
-class TP1DraftModelRunner(ModelRunnerWrapperBase):
41
- """Specialized model runner for speculative decoding draft model.
42
- Since the draft model always execute k forward passes consecutively to
43
- generate k speculative tokens in a single speculative decoding step,
44
- we could get rid of most CPU-GPU synchronization and data transfer
45
- overheads by keeping model input and output tensors on GPU all the time.
46
-
47
- TODOs:
48
- 1. Currently supports only flash-attn, add support for other attn_backends.
49
- 2. Support TP > 1 (this requires some designs because we do not expect
50
- any broadcasting inside execute_model).
51
- """
52
-
53
- def __init__(self, model_runner: ModelRunnerBase):
54
- super().__init__(model_runner)
55
-
56
- self.indices_of_seq_with_bonus_tokens = None
57
-
58
- def _update_sampling_metadata(self, sampling_metadata, num_seqs,
59
- num_queries):
60
-
61
- assert sampling_metadata.num_prompts == 0
62
- assert len(sampling_metadata.seq_groups) == num_queries
63
- assert sampling_metadata.selected_token_indices.shape == (
64
- num_queries, )
65
- # assert sampling_metadata.categorized_sample_indices == TODO: Add if needed # noqa: E501
66
-
67
- # Verify that all sequences are decodes
68
- for i in range(num_queries):
69
- seq_group = sampling_metadata.seq_groups[i]
70
-
71
- assert seq_group.is_prompt is False # No prompt
72
- assert seq_group.prompt_logprob_indices == [] # No prompt
73
- assert seq_group.sample_indices == [i] # Simple
74
-
75
- def _gpu_advance_step(self, model_input: ModelRunnerInputBase,
76
- last_output: SamplerOutput) -> ModelRunnerInputBase:
77
- # Currently, we expect "decode mode" only
78
- assert not model_input.is_prompt
79
-
80
- # Get num_seqs
81
- num_seqs = len(model_input.seq_lens)
82
- num_queries = len(model_input.query_lens)
83
-
84
- # Get output tokens GPU tensor
85
- sampled_token_ids = last_output.sampled_token_ids
86
- assert sampled_token_ids is not None
87
-
88
- # Update attn_metadata
89
- attn_metadata = model_input.attn_metadata
90
- assert isinstance(attn_metadata, FlashAttentionMetadata)
91
-
92
- attn_metadata.advance_step(model_input, sampled_token_ids,
93
- self.block_size, num_seqs, num_queries)
94
-
95
- # Update sampling_metadata
96
- sampling_metadata = model_input.sampling_metadata
97
- self._update_sampling_metadata(sampling_metadata, num_seqs,
98
- num_queries)
99
-
100
- # Create new input
101
- new_model_input = self._model_input_cls(
102
- input_tokens=model_input.input_tokens,
103
- input_positions=model_input.input_positions,
104
- attn_metadata=attn_metadata,
105
- seq_lens=attn_metadata.seq_lens,
106
- query_lens=model_input.query_lens,
107
- lora_mapping=model_input.lora_mapping,
108
- lora_requests=model_input.lora_requests,
109
- multi_modal_kwargs=model_input.multi_modal_kwargs,
110
- sampling_metadata=model_input.sampling_metadata,
111
- is_prompt=False,
112
- )
113
-
114
- # Ensure we skip CPU samples
115
- assert new_model_input.sampling_metadata.skip_sampler_cpu_output is True
116
- # We can reuse sampling tensors since every decode iteration is the same
117
- new_model_input.sampling_metadata.reuse_sampling_tensors = True
118
-
119
- if debug_advance_input:
120
- logger.debug("NEW INPUT: ")
121
- logger.debug(" input_tokens = %s", new_model_input.input_tokens)
122
- logger.debug(" input_positions = %s",
123
- new_model_input.input_positions)
124
- logger.debug(" seq_lens = %d", new_model_input.seq_lens)
125
- logger.debug(" query_lens = %d", new_model_input.query_lens)
126
- logger.debug(" attn_metadata:")
127
- logger.debug(" seq_lens_tensor: %s",
128
- attn_metadata.seq_lens_tensor)
129
- logger.debug(" slot_mapping: %s", attn_metadata.slot_mapping)
130
- logger.debug(" block_tables: %s", attn_metadata.block_tables)
131
-
132
- return new_model_input
133
-
134
- def supports_gpu_multi_step(self, execute_model_req: ExecuteModelRequest):
135
- """Determines if draft_model_runner GPU multi-step can be used.
136
- Currently required conditions are:
137
- 1. Only decodes
138
- 2. Only flash-attn
139
- 3. No LORA
140
- 4. No prompt_adapter_config
141
- """
142
- if not allow_gpu_advance_step:
143
- return False
144
-
145
- # We allow multi-step GPU only in decode mode
146
- for seq_group in execute_model_req.seq_group_metadata_list:
147
- if seq_group.is_prompt:
148
- return False
149
-
150
- # TODO: Add support for other attn backends
151
- if self.attn_backend.get_name() not in ("FLASH_ATTN", ):
152
- return False
153
-
154
- # TODO: Add support for LORA
155
- if self.lora_config:
156
- return False
157
-
158
- # TODO: Add soft-tuning prompt adapter support
159
- return not self.prompt_adapter_config
160
-
161
- def set_indices_of_seq_with_bonus_tokens(self,
162
- indices_of_seq_with_bonus_tokens):
163
- self.indices_of_seq_with_bonus_tokens = indices_of_seq_with_bonus_tokens
164
-
165
- @torch.inference_mode()
166
- def execute_model(
167
- self,
168
- model_input: ModelRunnerInputBase,
169
- kv_caches: List[torch.Tensor],
170
- previous_hidden_states: Optional[torch.Tensor] = None,
171
- intermediate_tensors: Optional[IntermediateTensors] = None,
172
- num_steps: int = 1,
173
- **kwargs,
174
- ) -> Optional[List[SamplerOutput]]:
175
- """Executes num_steps forward passes with advacement of input tensors
176
- on the GPU. Look at supports_gpu_multi_step(..) for pre-conditions.
177
-
178
- Optimizations used:
179
- 1. Input tensors are updated on the GPU directly
180
- 2. Skips GPU=>CPU serialization of sampler outputs (we don't need
181
- them since we do batch expansion later that uses GPU outputs)
182
- 3. Reuses sampling tensors (since we run only decodes and they have
183
- a repeating sampling logic)
184
- """
185
-
186
- # When num_steps == 1, we execute the fallback here for the GPU
187
- # advance_step, which runs prepare_inputs on CPU and for each spec
188
- # iteration invokes this function only once
189
- # (Look at multi-step-worker code)
190
- is_fallback = num_steps == 1
191
- if not is_fallback:
192
- # Since we do not broadcast data inside execute_model anymore,
193
- # we need to figure out the best way to support TP > 1 in this
194
- # case, because we will at least need to broadcast the sampled
195
- # tokens to all workers.
196
- if not self.is_driver_worker:
197
- raise ValueError("TP1DraftModelRunner only supports TP=1.")
198
-
199
- # Sanity
200
- if self.lora_config is not None:
201
- raise ValueError("TP1DraftModelRunner has no support for LORA")
202
- if self.prompt_adapter_config is not None:
203
- raise ValueError("TP1DraftModelRunner has no support for "
204
- "prompt_adapter_config")
205
- if model_input.inputs_embeds is not None:
206
- raise ValueError("TP1DraftModelRunner has no support for "
207
- "inputs_embeds")
208
- if model_input.multi_modal_kwargs:
209
- raise ValueError(
210
- "TP1DraftModelRunner has no support for multi_modal_kwargs"
211
- )
212
- else:
213
- if self.lora_config:
214
- assert model_input.lora_requests is not None
215
- assert model_input.lora_mapping is not None
216
- self.set_active_loras(model_input.lora_requests,
217
- model_input.lora_mapping)
218
-
219
- if self.prompt_adapter_config:
220
- assert model_input.prompt_adapter_requests is not None
221
- assert model_input.prompt_adapter_mapping is not None
222
- self.set_active_prompt_adapters(
223
- model_input.prompt_adapter_requests,
224
- model_input.prompt_adapter_mapping)
225
-
226
- self.attn_state.begin_forward(model_input)
227
-
228
- # Detect exec mode
229
- assert model_input.attn_metadata is not None
230
- use_cuda_graph = False
231
- if model_input.attn_metadata.num_prefills > 0:
232
- # In this case, execute_model(..) was called directly
233
- if num_steps > 1:
234
- raise ValueError(
235
- "execute_model(..) of draft_model_runner can be called "
236
- "directly only with a single-step prefill")
237
- else:
238
- # We can skip CPU samples for spec token generation.
239
- # (We do allow CPU samples for num_steps == 1 to support the
240
- # fallback case, where supports_gpu_multi_step(..) does not pass)
241
- model_input.sampling_metadata.skip_sampler_cpu_output = (
242
- not is_fallback)
243
-
244
- # Attn attr defines if we use cuda graphs
245
- use_cuda_graph = model_input.attn_metadata.use_cuda_graph
246
-
247
- # Get model
248
- if use_cuda_graph:
249
- if model_input.inputs_embeds is None:
250
- graph_batch_size = model_input.input_tokens.shape[0]
251
- model_executable = (
252
- self.graph_runners[model_input.virtual_engine][(
253
- graph_batch_size, False)])
254
- else:
255
- graph_batch_size = model_input.inputs_embeds.shape[0]
256
- model_executable = (
257
- self.graph_runners[model_input.virtual_engine][(
258
- graph_batch_size, True)])
259
-
260
- if previous_hidden_states is not None:
261
- hidden_states = torch.cat([
262
- previous_hidden_states,
263
- torch.empty([
264
- graph_batch_size - previous_hidden_states.shape[0],
265
- *previous_hidden_states.shape[1:]
266
- ],
267
- dtype=previous_hidden_states.dtype,
268
- device=previous_hidden_states.device)
269
- ])
270
- else:
271
- hidden_states = None
272
- else:
273
- model_executable = self.model
274
- hidden_states = previous_hidden_states
275
-
276
- outputs: List[SamplerOutput] = []
277
- for step in range(num_steps):
278
- multi_modal_kwargs = model_input.multi_modal_kwargs or {}
279
-
280
- model_execute_kwargs = {"previous_hidden_states": hidden_states} \
281
- if previous_hidden_states is not None else {}
282
-
283
- compute_logits_kwargs = {}
284
- # Run model
285
- if hasattr(self.model.config, "num_nextn_predict_layers"):
286
- # for DeepSeek MTP only to use the corresponding layer for
287
- # each step
288
- spec_step_idx = kwargs.get("spec_step_idx", step)
289
- model_execute_kwargs["spec_step_idx"] = spec_step_idx
290
- compute_logits_kwargs["spec_step_idx"] = spec_step_idx
291
- with set_forward_context(model_input.attn_metadata,
292
- self.vllm_config):
293
- hidden_states = model_executable(
294
- input_ids=model_input.input_tokens,
295
- inputs_embeds=None,
296
- positions=model_input.input_positions,
297
- intermediate_tensors=intermediate_tensors,
298
- **MultiModalKwargs.as_kwargs(
299
- multi_modal_kwargs,
300
- device=self.device,
301
- ),
302
- **model_execute_kwargs,
303
- )
304
-
305
- # Compute the logits.
306
- logits = self.model.compute_logits(hidden_states,
307
- model_input.sampling_metadata,
308
- **compute_logits_kwargs)
309
- if not self.is_driver_worker:
310
- return []
311
- # Sample the next token.
312
- output = self.model_runner.sampler(
313
- logits=logits,
314
- sampling_metadata=model_input.sampling_metadata,
315
- )
316
- outputs.append(output)
317
-
318
- if self.return_hidden_states and is_fallback:
319
- if use_cuda_graph:
320
- indices = model_input.sampling_metadata\
321
- .selected_token_indices
322
- output.hidden_states = hidden_states[:len(indices)]
323
- else:
324
- output.hidden_states = hidden_states
325
-
326
- if model_input.attn_metadata.num_prefills == 0 \
327
- and self.indices_of_seq_with_bonus_tokens is not None:
328
- assert output.sampled_token_ids is not None
329
- # output.sampled_token_ids should be of shape (num_seqs, 1)
330
- nums_seqs, num_tokens_per_seq = output.sampled_token_ids.shape
331
- assert num_tokens_per_seq == 1
332
- count = 0
333
- for i in range(nums_seqs):
334
- bonus_seq_idx = self.indices_of_seq_with_bonus_tokens[
335
- count]
336
- if i != bonus_seq_idx:
337
- # The following might cause a cpu->gpu sync
338
- # However, the performance impact is negligible as we
339
- # benchmarked on H100.
340
- output.sampled_token_ids[
341
- i, :] = model_input.input_tokens[bonus_seq_idx]
342
- else:
343
- count += 1
344
-
345
- # Prepare inputs for the next step
346
- if step != num_steps - 1:
347
- model_input = self._gpu_advance_step(model_input, outputs[-1])
348
-
349
- return outputs
vllm/spec_decode/interfaces.py
deleted
-99
@@ -1,99 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from abc import ABC, abstractmethod
5
-from dataclasses import dataclass
6
-from typing import List, Optional, Set, Union
7
-
8
-import torch
9
-
10
-from vllm.sequence import ExecuteModelRequest, PromptLogprobs
11
-from vllm.worker.worker_base import WorkerBase
12
-
13
-
14
-@dataclass
15
-class SpeculativeProposals:
16
- """Datastructure used to represent proposal tokens from some proposer. It
17
- also tracks how many speculative tokens each sequence has.
18
- """
19
-
20
- # Speculative proposal tokens.
21
- proposal_token_ids: torch.Tensor
22
-
23
- # Probabilities of the proposal tokens according to the proposer.
24
- proposal_probs: torch.Tensor
25
-
26
- # The valid length of each proposal; can be zero.
27
- proposal_lens: torch.Tensor
28
-
29
- # A flag to mark that there's no available proposals
30
- no_proposals: bool = False
31
-
32
- def __repr__(self):
33
- return (f"SpeculativeProposals("
34
- f"proposal_token_ids={self.proposal_token_ids}, "
35
- f"proposal_probs={self.proposal_probs.shape}, "
36
- f"proposal_lens={self.proposal_lens})")
37
-
38
-
39
-@dataclass
40
-class SpeculativeScores:
41
- """Datastructure used to represent the scores of speculative tokens
42
- according to the scoring model.
43
- """
44
-
45
- # Probabilities of the speculative tokens according to the scoring model.
46
- probs: torch.Tensor
47
-
48
- # Log-probabilities of the speculative tokens according to the scoring
49
- # model. These values can be used to generate Logprob objects that are
50
- # returned to the user.
51
- logprobs: torch.Tensor
52
-
53
- # Token ids sampled from the scoring model. Used for speculative bonus
54
- # tokens and also non-speculative normal decoding.
55
- token_ids: torch.Tensor
56
-
57
- # Optional last hidden states from the scoring model.
58
- hidden_states: Optional[torch.Tensor] = None
59
-
60
- # Scoring model may also return logprobs for prompt tokens
61
- # for each request, when chunked prefill is enabled.
62
- prompt_logprobs: Optional[List[PromptLogprobs]] = None
63
-
64
- def __repr__(self):
65
- return (f"SpeculativeScores("
66
- f"probs={self.probs.shape}, "
67
- f"token_ids={self.token_ids.shape})")
68
-
69
-
70
-class SpeculativeProposer(ABC):
71
-
72
- @abstractmethod
73
- def get_spec_proposals(
74
- self,
75
- execute_model_req: ExecuteModelRequest,
76
- # If set, this contains all sequence IDs that were assigned
77
- # bonus tokens in their last forward pass.
78
- seq_ids_with_bonus_token_in_last_step: Set[int],
79
- ) -> SpeculativeProposals:
80
- raise NotImplementedError
81
-
82
-
83
-class SpeculativeScorer(ABC):
84
-
85
- def __init__(self, scorer_worker: WorkerBase,
86
- device: Union[torch.device, str], vocab_size: int):
87
- self._scorer_worker = scorer_worker
88
- if isinstance(device, torch.device):
89
- device = device.type
90
- self._device = device
91
- self._vocab_size = vocab_size
92
-
93
- @abstractmethod
94
- def score_proposals(
95
- self,
96
- execute_model_req: ExecuteModelRequest,
97
- proposals: SpeculativeProposals,
98
- ) -> SpeculativeScores:
99
- raise NotImplementedError
vllm/spec_decode/medusa_worker.py
deleted
-138
@@ -1,138 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import weakref
5
-from typing import List, Optional, Set, Tuple
6
-
7
-import torch
8
-
9
-from vllm.model_executor import SamplingMetadata
10
-from vllm.model_executor.layers.sampler import SamplerOutput
11
-from vllm.sequence import ExecuteModelRequest, SequenceGroupMetadata
12
-from vllm.spec_decode.interfaces import SpeculativeProposals
13
-from vllm.spec_decode.proposer_worker_base import NonLLMProposerWorkerBase
14
-from vllm.spec_decode.top1_proposer import Top1Proposer
15
-from vllm.worker.worker_base import DelegateWorkerBase
16
-
17
-
18
-class MedusaWorker(NonLLMProposerWorkerBase, DelegateWorkerBase):
19
- """Worker for Medusa.
20
- """
21
-
22
- def __init__(self, *args, **kwargs):
23
- DelegateWorkerBase.__init__(self, *args, **kwargs)
24
- # Lazy initialization list.
25
- self._proposer: Top1Proposer
26
-
27
- def init_device(self):
28
- self.worker.init_device()
29
-
30
- self._proposer = Top1Proposer(
31
- weakref.proxy(self), # type: ignore[arg-type]
32
- self.device,
33
- self.vocab_size,
34
- max_proposal_len=self.max_model_len,
35
- )
36
-
37
- def set_include_gpu_probs_tensor(self):
38
- pass
39
-
40
- def set_should_modify_greedy_probs_inplace(self):
41
- pass
42
-
43
- @torch.inference_mode()
44
- def sampler_output(
45
- self,
46
- execute_model_req: ExecuteModelRequest,
47
- sample_len: int,
48
- # Unused parameter.
49
- seq_ids_with_bonus_token_in_last_step: Set[int],
50
- ) -> Tuple[List[SamplerOutput], bool]:
51
- """Run the model forward pass to generate sample_len future tokens.
52
- Returns the list of sampler output, one per layer, along with indicator
53
- of whether torch tensor in sampler output need to be transposed in
54
- latter sampler_output_to_torch logic.
55
-
56
- For medusa worker, this indicator shall be False.
57
- """
58
- self._raise_if_unsupported(execute_model_req)
59
-
60
- seq_group_metadata_list = execute_model_req.seq_group_metadata_list
61
-
62
- seq_lens, query_lens = self._prepare_input_tensors(
63
- seq_group_metadata_list)
64
-
65
- generators = self.model_runner.get_generators(
66
- execute_model_req.finished_requests_ids)
67
- sampling_metadata = SamplingMetadata.prepare(
68
- seq_group_metadata_list, seq_lens, query_lens, self.device,
69
- self.model_runner.pin_memory, generators)
70
-
71
- model_outputs = self.model_runner.model.generate_proposals(
72
- previous_hidden_states=execute_model_req.previous_hidden_states.
73
- hidden_states,
74
- sampling_metadata=sampling_metadata)
75
-
76
- return model_outputs, False
77
-
78
- def _prepare_input_tensors(
79
- self,
80
- seq_group_metadata_list: Optional[List[SequenceGroupMetadata]],
81
- ) -> Tuple[List[int], List[int]]:
82
- if not seq_group_metadata_list:
83
- return [], []
84
-
85
- seq_lens: List[int] = []
86
- query_lens: List[int] = []
87
-
88
- for seq_group_metadata in seq_group_metadata_list:
89
- is_prompt = seq_group_metadata.is_prompt
90
-
91
- for seq_data in seq_group_metadata.seq_data.values():
92
- seq_data_len = seq_data.get_len()
93
- if is_prompt:
94
- context_len = seq_data.get_num_computed_tokens()
95
- seq_len = min(
96
- seq_data_len,
97
- context_len + seq_group_metadata.token_chunk_size)
98
- seq_lens.append(seq_len)
99
- query_lens.append(seq_len - context_len)
100
- else:
101
- seq_lens.append(seq_data_len)
102
- query_lens.append(1)
103
-
104
- return seq_lens, query_lens
105
-
106
- def get_spec_proposals(
107
- self,
108
- execute_model_req: ExecuteModelRequest,
109
- seq_ids_with_bonus_token_in_last_step: Set[int],
110
- ) -> SpeculativeProposals:
111
- """Produce speculations given an input batch of sequences. The number of
112
- speculative tokens per sequence is determined by max_proposal_len.
113
- """
114
-
115
- return self._proposer.get_spec_proposals(
116
- execute_model_req, seq_ids_with_bonus_token_in_last_step)
117
-
118
- def _raise_if_unsupported(
119
- self,
120
- execute_model_req: ExecuteModelRequest,
121
- ) -> None:
122
- """MedusaWorker does not yet implement support for cache swap
123
- operations or beam search.
124
- """
125
- if any([
126
- execute_model_req.blocks_to_swap_in,
127
- execute_model_req.blocks_to_swap_out,
128
- execute_model_req.blocks_to_copy
129
- ]):
130
- raise NotImplementedError(
131
- "MedusaWorker does not support cache operations")
132
-
133
- if any(
134
- len(seq_group_metadata.seq_data.keys()) != 1
135
- for seq_group_metadata in
136
- execute_model_req.seq_group_metadata_list):
137
- raise NotImplementedError(
138
- "MedusaWorker does not support beam search.")
vllm/spec_decode/metrics.py
deleted
-213
@@ -1,213 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import time
5
-from typing import Callable, Optional, Union
6
-
7
-import msgspec
8
-import torch
9
-
10
-from vllm.model_executor.layers.spec_decode_base_sampler import (
11
- SpecDecodeBaseSampler)
12
-from vllm.platforms import current_platform
13
-from vllm.utils import is_pin_memory_available
14
-
15
-
16
-class SpecDecodeWorkerMetrics(
17
- msgspec.Struct,
18
- omit_defaults=True, # type: ignore[call-arg]
19
- array_like=True): # type: ignore[call-arg]
20
- """Dataclass holding metrics emitted from the spec decode worker.
21
- """
22
-
23
- # The empirical acceptance rate of the proposal method on a per-token basis.
24
- # This is useful for evaluating how well the proposal method aligns with the
25
- # scoring method.
26
- draft_acceptance_rate: float
27
-
28
- # The empirical efficiency, measured as the number of tokens emitted by the
29
- # system divided by the number of tokens that could be emitted by the system
30
- # if the proposal method were perfect.
31
- system_efficiency: float
32
-
33
- # The number of speculative tokens produced by the proposal method.
34
- draft_tokens: int
35
-
36
- # The number of tokens emitted by the entire system.
37
- emitted_tokens: int
38
-
39
- # The number of tokens accepted by the scoring model and verification
40
- # routine, e.g. Llama2-70B and lossless rejection sampling.
41
- #
42
- # NOTE: Any token accepted by the verification routine is considered
43
- # accepted (regardless of if the speculative prefix is also accepted). The
44
- # user will usually see less accepted tokens. This metric is helpful when
45
- # evaluating alignment of the proposal method with the scoring model.
46
- accepted_tokens: int
47
-
48
- # The number of speculative tokens per sequence.
49
- num_spec_tokens: int
50
-
51
-
52
-Timer = Callable[[], float]
53
-
54
-
55
-class AsyncMetricsCollector:
56
- """Class which copies rejection/typical-acceptance sampler metrics
57
- from the device to CPU on a non-default Torch stream.
58
- """
59
-
60
- def __init__(self,
61
- spec_decode_sampler: SpecDecodeBaseSampler,
62
- timer: Optional[Timer] = None,
63
- collect_interval_s: float = 5.0):
64
- self.spec_decode_sampler = spec_decode_sampler
65
- self._timer = time.time if timer is None else timer
66
-
67
- self._rank: Optional[int] = None
68
-
69
- # We don't have a device set yet.
70
- self._copy_stream: Optional[torch.cuda.Stream] = None
71
-
72
- self._in_flight_copy: Optional[torch.cuda.Event] = None
73
-
74
- pin_memory = is_pin_memory_available()
75
- self._aggregate_num_accepted_tokens = torch.tensor(
76
- 0, dtype=torch.long, device="cpu", pin_memory=pin_memory)
77
- self._aggregate_num_emitted_tokens = torch.tensor(
78
- 0, dtype=torch.long, device="cpu", pin_memory=pin_memory)
79
- self._aggregate_num_draft_tokens = 0
80
-
81
- self._rejsample_metrics_collect_interval_s = collect_interval_s
82
- self._last_metrics_collect_time = self._timer()
83
-
84
- def init_gpu_tensors(self, rank: int) -> None:
85
- self._rank = rank
86
- self._copy_stream = torch.cuda.Stream()
87
-
88
- def init_tensors(self,
89
- rank: int,
90
- device_type: Union[torch.device, str] = 'cuda') -> None:
91
- self._rank = rank
92
- if isinstance(device_type, torch.device):
93
- device_type = device_type.type
94
- stream = current_platform.Stream
95
- if stream is not None:
96
- self._copy_stream = stream()
97
-
98
- def maybe_collect_rejsample_metrics(
99
- self, k: int) -> Optional[SpecDecodeWorkerMetrics]:
100
- # Skip for any platform that doesn't have device Event
101
- if current_platform.Event is None:
102
- return None
103
-
104
- # If a copy was initiated in the previous call, collect and return.
105
- if self._in_flight_copy is not None:
106
- ready_event = self._in_flight_copy
107
- self._in_flight_copy = None
108
- return self._collect_rejsample_metrics(k, ready_event)
109
-
110
- # Otherwise, check if we should start a new copy.
111
- if self._should_collect_rejsample_metrics(self._timer()):
112
- assert self._in_flight_copy is None
113
- self._in_flight_copy = self._copy_rejsample_metrics_async()
114
-
115
- return None
116
-
117
- def _should_collect_rejsample_metrics(self, now: float) -> bool:
118
- """Return whether or not this iteration should print sampling
119
- metrics.
120
- """
121
- if self._rank != 0:
122
- return False
123
-
124
- return now - self._last_metrics_collect_time >= self._rejsample_metrics_collect_interval_s # noqa: E501
125
-
126
- def _copy_rejsample_metrics_async(self) -> torch.cuda.Event:
127
- """Copy rejection/typical-acceptance sampling metrics
128
- (number of accepted tokens, etc) to CPU asynchronously.
129
-
130
- Returns a device event recording when the copy is complete.
131
- """
132
- assert self._copy_stream is not None
133
- self._copy_stream.wait_stream(current_platform.current_stream())
134
-
135
- with current_platform.stream(self._copy_stream):
136
- self._aggregate_num_accepted_tokens.copy_(
137
- self.spec_decode_sampler.num_accepted_tokens,
138
- non_blocking=True)
139
- self._aggregate_num_emitted_tokens.copy_(
140
- self.spec_decode_sampler.num_emitted_tokens, non_blocking=True)
141
- # Number of draft tokens is calculated on CPU, so no copy is
142
- # required.
143
- self._aggregate_num_draft_tokens = (
144
- self.spec_decode_sampler.num_draft_tokens)
145
-
146
- aggregate_metrics_ready = current_platform.Event()
147
- aggregate_metrics_ready.record(self._copy_stream)
148
-
149
- return aggregate_metrics_ready
150
-
151
- def _collect_rejsample_metrics(
152
- self, k: int,
153
- ready_event: torch.cuda.Event) -> SpecDecodeWorkerMetrics:
154
- """Create metrics object from statistics copied asynchronously.
155
-
156
- Args:
157
- k: int. The number of speculative tokens; used to determine system
158
- efficiency.
159
- ready_event: torch.cuda.Event. The CUDA event recording when the
160
- async GPU->CPU copy is complete.
161
- """
162
-
163
- ready_event.synchronize()
164
-
165
- # update time of last collection
166
- self._last_metrics_collect_time = self._timer()
167
-
168
- accepted_tokens = self._aggregate_num_accepted_tokens.item()
169
- emitted_tokens = self._aggregate_num_emitted_tokens.item()
170
- draft_tokens = self._aggregate_num_draft_tokens
171
-
172
- max_num_emitted_tokens = self.get_max_num_emitted_tokens(
173
- draft_tokens, k)
174
-
175
- if draft_tokens > 0:
176
- draft_acceptance_rate = accepted_tokens / draft_tokens
177
- else:
178
- draft_acceptance_rate = float("nan")
179
-
180
- if max_num_emitted_tokens > 0:
181
- system_efficiency = emitted_tokens / max_num_emitted_tokens
182
- else:
183
- system_efficiency = float("nan")
184
-
185
- return SpecDecodeWorkerMetrics(
186
- num_spec_tokens=k,
187
- draft_acceptance_rate=draft_acceptance_rate,
188
- system_efficiency=system_efficiency,
189
- accepted_tokens=accepted_tokens,
190
- draft_tokens=draft_tokens,
191
- emitted_tokens=emitted_tokens,
192
- )
193
-
194
- @staticmethod
195
- def get_max_num_emitted_tokens(draft_tokens: int, k: int) -> int:
196
- """Calculate the number of emitted tokens, assuming all tokens are
197
- accepted.
198
-
199
- This is equal to the number of sequences that have been speculated on,
200
- times (speculation len + 1). The +1 comes from the bonus token.
201
- """
202
- # Determine the number of sequences that have been speculated on. Since
203
- # the batch size can be variable, we divide by k.
204
- assert draft_tokens % k == 0
205
- total_num_spec_seqs = draft_tokens // k
206
-
207
- # A single sequence may emit k accepted tokens and one bonus token in
208
- # the best case.
209
- num_emitted_per_seq_if_all_accepted = k + 1
210
-
211
- # The max num of emitted tokens is the number of speculated sequences
212
- # times the max emitted per seq.
213
- return total_num_spec_seqs * num_emitted_per_seq_if_all_accepted
vllm/spec_decode/mlp_speculator_worker.py
deleted
-94
@@ -1,94 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from typing import List, Optional, Set, Tuple
5
-
6
-import torch
7
-
8
-from vllm.model_executor import SamplingMetadata
9
-from vllm.model_executor.layers.sampler import SamplerOutput
10
-from vllm.sequence import ExecuteModelRequest, SequenceGroupMetadata
11
-from vllm.spec_decode.multi_step_worker import MultiStepWorker
12
-from vllm.spec_decode.proposer_worker_base import NonLLMProposerWorkerBase
13
-
14
-
15
-class MLPSpeculatorWorker(NonLLMProposerWorkerBase, MultiStepWorker):
16
- """Worker for MLPSpeculator models.
17
-
18
- Not currently compatible with LoRA or chunked prefill.
19
- """
20
-
21
- @torch.inference_mode()
22
- def sampler_output(
23
- self,
24
- execute_model_req: ExecuteModelRequest,
25
- sample_len: int,
26
- # Unused parameter. MLPSpeculatorWorker does not use the KV Cache and
27
- # therefore does not need this parameter.
28
- seq_ids_with_bonus_token_in_last_step: Set[int],
29
- ) -> Tuple[List[SamplerOutput], bool]:
30
- """Run the model forward pass to generate sample_len future tokens.
31
- Returns the list of sampler output, one per layer, along with indicator
32
- of whether torch tensor in sampler output need to be transposed in
33
- latter sampler_output_to_torch logic.
34
-
35
- For mlp spec worker, this indicator shall be True.
36
- """
37
- self._raise_if_unsupported(execute_model_req)
38
-
39
- seq_group_metadata_list = execute_model_req.seq_group_metadata_list
40
-
41
- (input_tokens, seq_lens,
42
- query_lens) = self._prepare_input_tensors(seq_group_metadata_list)
43
-
44
- generators = self.model_runner.get_generators(
45
- execute_model_req.finished_requests_ids)
46
- sampling_metadata = SamplingMetadata.prepare(
47
- seq_group_metadata_list, seq_lens, query_lens, self.device,
48
- self.model_runner.pin_memory, generators)
49
-
50
- model_outputs = self.model_runner.model.generate_proposals(
51
- input_ids=input_tokens,
52
- previous_hidden_states=execute_model_req.previous_hidden_states.
53
- hidden_states,
54
- num_predict_tokens=sample_len,
55
- sampling_metadata=sampling_metadata)
56
-
57
- assert len(model_outputs) == sample_len
58
-
59
- return model_outputs, True
60
-
61
- def _prepare_input_tensors(
62
- self,
63
- seq_group_metadata_list: Optional[List[SequenceGroupMetadata]],
64
- ) -> Tuple[torch.Tensor, List[int], List[int]]:
65
- if not seq_group_metadata_list:
66
- return torch.empty(0, device=self.device), [], []
67
-
68
- input_tokens: List[int] = []
69
- seq_lens: List[int] = []
70
- query_lens: List[int] = []
71
-
72
- for seq_group_metadata in seq_group_metadata_list:
73
- is_prompt = seq_group_metadata.is_prompt
74
-
75
- for seq_data in seq_group_metadata.seq_data.values():
76
- seq_data_len = seq_data.get_len()
77
- if is_prompt:
78
- context_len = seq_data.get_num_computed_tokens()
79
- seq_len = min(
80
- seq_data_len,
81
- context_len + seq_group_metadata.token_chunk_size)
82
- tokens = seq_data.get_token_ids()[context_len:seq_len]
83
- seq_lens.append(seq_len)
84
- input_tokens.extend(tokens)
85
- query_lens.append(seq_len - context_len)
86
- else:
87
- seq_lens.append(seq_data_len)
88
- input_tokens.append(seq_data.get_last_token_id())
89
- query_lens.append(1)
90
-
91
- input_tokens_tensor = torch.tensor(input_tokens,
92
- dtype=torch.long,
93
- device=self.device)
94
- return input_tokens_tensor, seq_lens, query_lens
vllm/spec_decode/mqa_scorer.py
deleted
-160
@@ -1,160 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from vllm.sequence import (ExecuteModelRequest, SequenceData,
5
- SequenceGroupMetadata, get_all_seq_ids)
6
-from vllm.spec_decode.interfaces import (SpeculativeProposals,
7
- SpeculativeScorer, SpeculativeScores)
8
-
9
-SeqId = int
10
-TargetSeqId = int
11
-
12
-
13
-class MQAScorer(SpeculativeScorer):
14
-
15
- def score_proposals(
16
- self,
17
- execute_model_req: ExecuteModelRequest,
18
- proposals: SpeculativeProposals,
19
- ) -> SpeculativeScores:
20
- target_seq_group_metadata_list = []
21
- target_seq_id_start = max(
22
- get_all_seq_ids(execute_model_req.seq_group_metadata_list)) + 1
23
- all_proposal_tokens = proposals.proposal_token_ids.tolist()
24
- all_proposal_lengths = proposals.proposal_lens.tolist()
25
- for i, seq_group_metadata in enumerate(
26
- execute_model_req.seq_group_metadata_list):
27
- if all_proposal_lengths[i] == 0:
28
- # Keep prompt seqs untouched (keep computed_tokens for chunks).
29
- target_seq_group_metadata_list.append(seq_group_metadata)
30
- continue
31
-
32
- seq_data_dict = seq_group_metadata.seq_data
33
- assert len(seq_data_dict) == 1
34
- seq_id = next(iter(seq_data_dict.keys()))
35
-
36
- seq_data: SequenceData = seq_data_dict[seq_id]
37
- prompt_token_ids = seq_data.get_prompt_token_ids()
38
- output_token_ids = seq_data.get_output_token_ids()
39
- proposal_token_ids = all_proposal_tokens[
40
- i][:all_proposal_lengths[i]]
41
- new_output_token_ids = [*output_token_ids, *proposal_token_ids]
42
-
43
- target_seq_id = target_seq_id_start + i
44
- new_seq_data = SequenceData.from_seqs(
45
- prompt_token_ids=prompt_token_ids,
46
- output_token_ids=new_output_token_ids,
47
- )
48
- new_seq_data.update_num_computed_tokens(
49
- len(prompt_token_ids) + len(output_token_ids) - 1)
50
-
51
- # Ensure that the new decode sequence has at least one token.
52
- assert len(output_token_ids) >= 1
53
- new_seq_data_dict = {target_seq_id: new_seq_data}
54
-
55
- new_seq_group_metadata = SequenceGroupMetadata(
56
- request_id=seq_group_metadata.request_id,
57
- is_prompt=seq_group_metadata.is_prompt,
58
- seq_data=new_seq_data_dict,
59
- sampling_params=seq_group_metadata.sampling_params,
60
- block_tables={
61
- target_seq_id: seq_group_metadata.block_tables[seq_id],
62
- },
63
- lora_request=None,
64
- )
65
- target_seq_group_metadata_list.append(new_seq_group_metadata)
66
-
67
- target_sampler_output = self._scorer_worker.execute_model(
68
- execute_model_req=execute_model_req.clone(
69
- seq_group_metadata_list=target_seq_group_metadata_list))
70
-
71
- target_sampler_output = target_sampler_output[0]
72
-
73
- k = execute_model_req.num_lookahead_slots
74
- bs = len(execute_model_req.seq_group_metadata_list)
75
- target_token_ids = target_sampler_output.sampled_token_ids
76
- target_probs = target_sampler_output.sampled_token_probs
77
- target_logprobs = target_sampler_output.logprobs
78
- prompt_logprobs = None
79
-
80
- # If all requests have the same number of query tokens, we can avoid
81
- # the for loop to build output for better performance.
82
- if min(all_proposal_lengths) == k:
83
- # Regular decodes only.
84
- assert all(not sg.is_prompt
85
- for sg in target_seq_group_metadata_list
86
- if sg.is_prompt)
87
- bs, _ = proposals.proposal_token_ids.shape
88
- all_tokens = target_token_ids.reshape(bs, k + 1)
89
- all_probs = target_probs.reshape(bs, k + 1, self._vocab_size)
90
- all_logprobs = target_logprobs.reshape(bs, k + 1, self._vocab_size)
91
- else:
92
- # We either have decodes with different lens or prefill+decodes.
93
- all_tokens = target_token_ids.new_full(size=(bs, k + 1),
94
- fill_value=-1)
95
- all_probs = target_probs.new_zeros(*all_tokens.shape,
96
- self._vocab_size)
97
- all_logprobs = target_logprobs.new_full(size=all_probs.shape,
98
- fill_value=-float("inf"))
99
- target_token_ids = target_token_ids.flatten()
100
-
101
- # When prompt logprobs is enabled, lens of returned tensors go from
102
- # n_sampled (requests with do_sample=True) to n_prompt+n_prefills.
103
- # We adjust stride accordingly to get the generated tokens and
104
- # their probs, but pass on prompt_logprobs as is, since it may be
105
- # that n_prompts >> K.
106
- has_prompt_log = any((sg.sampling_params.prompt_logprobs
107
- and sg.sampling_params.prompt_logprobs > 0)
108
- for sg in target_seq_group_metadata_list)
109
- # TODO (NickLucche) we should surface `disable_logprobs` as to not
110
- # break abstraction to get its value.
111
- if (not self._scorer_worker.model_runner.disable_logprobs\
112
- and has_prompt_log):
113
- prompt_logprobs = [
114
- o.prompt_logprobs for o in target_sampler_output.outputs
115
- ]
116
-
117
- # Split loop into prefill|decode for readability.
118
- start_loc, i = 0, 0
119
- while i < len(target_seq_group_metadata_list
120
- ) and target_seq_group_metadata_list[i].is_prompt:
121
- seq_meta = target_seq_group_metadata_list[i]
122
- end_loc = start_loc
123
- if has_prompt_log:
124
- end_loc += seq_meta.token_chunk_size
125
- elif seq_meta.do_sample:
126
- end_loc += 1
127
-
128
- # Skip chunks with no output tokens.
129
- if seq_meta.do_sample:
130
- # Get sampled token (last position in chunk) and its prob.
131
- all_tokens[i, 0] = target_token_ids[end_loc - 1]
132
- all_probs[i, 0] = target_probs[end_loc - 1]
133
- all_logprobs[i, 0] = target_logprobs[end_loc - 1]
134
-
135
- i += 1
136
- start_loc = end_loc
137
- # Decodes.
138
- while i < len(target_seq_group_metadata_list):
139
- proposed_len, seq_meta = all_proposal_lengths[
140
- i], target_seq_group_metadata_list[i]
141
- output_len = proposed_len + 1
142
- end_loc = start_loc + output_len
143
- all_tokens[
144
- i, :output_len] = target_token_ids[start_loc:end_loc]
145
- all_probs[i, :output_len] = target_probs[start_loc:end_loc]
146
- all_logprobs[
147
- i, :output_len] = target_logprobs[start_loc:end_loc]
148
- start_loc = end_loc
149
- i += 1
150
-
151
- hidden_states = None
152
- if target_sampler_output.hidden_states is not None:
153
- hidden_states = target_sampler_output.hidden_states.reshape(
154
- bs, (k + 1), -1)
155
-
156
- return SpeculativeScores(probs=all_probs,
157
- token_ids=all_tokens,
158
- logprobs=all_logprobs,
159
- hidden_states=hidden_states,
160
- prompt_logprobs=prompt_logprobs)
vllm/spec_decode/multi_step_worker.py
deleted
-423
@@ -1,423 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import copy
5
-import weakref
6
-from typing import Dict, List, Set, Tuple
7
-
8
-import torch
9
-
10
-from vllm.model_executor.layers.sampler import SamplerOutput
11
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
12
-from vllm.platforms import current_platform
13
-from vllm.sequence import (ExecuteModelRequest, HiddenStates, SequenceData,
14
- SequenceGroupMetadata)
15
-
16
-if current_platform.is_cuda_alike():
17
- from vllm.spec_decode.draft_model_runner import TP1DraftModelRunner
18
-
19
-from vllm.spec_decode.interfaces import (SpeculativeProposals,
20
- SpeculativeProposer)
21
-from vllm.spec_decode.proposer_worker_base import ProposerWorkerBase
22
-from vllm.spec_decode.top1_proposer import Top1Proposer
23
-from vllm.worker.worker_base import DelegateWorkerBase
24
-
25
-
26
-class MultiStepWorker(ProposerWorkerBase, DelegateWorkerBase):
27
- """The MultiStepWorker is equivalent to a Worker except that it allows
28
- multiple forward passes in a single call, assuming the scheduler has
29
- allocated enough space to store the additional KV. This reduces overhead
30
- by invoking the scheduler less.
31
-
32
- The MultiStepWorker does not support cache swap operations, or beam search.
33
- Cache swap operations do not require large modifications. On the other hand,
34
- beam search requires memory allocations during sequence forks and thus
35
- requires more thought for MultiStepWorker support.
36
- """
37
-
38
- def __init__(self, *args, **kwargs):
39
- DelegateWorkerBase.__init__(self, *args, **kwargs)
40
- # Lazy initialization list.
41
- self._proposer: SpeculativeProposer
42
-
43
- def init_device(self) -> None:
44
- self.worker.init_device()
45
- self._proposer = Top1Proposer(
46
- weakref.proxy(self), # type: ignore[arg-type]
47
- self.device,
48
- self.vocab_size,
49
- max_proposal_len=self.max_model_len,
50
- )
51
-
52
- def set_include_gpu_probs_tensor(self) -> None:
53
- # Need include_gpu_probs_tensor for MultiStepWorker
54
- self.model_runner.sampler.include_gpu_probs_tensor = True
55
- if hasattr(self.model_runner.model, "sampler"):
56
- (self.model_runner.model.sampler.include_gpu_probs_tensor) = True
57
-
58
- def set_should_modify_greedy_probs_inplace(self) -> None:
59
- self.model_runner.sampler.should_modify_greedy_probs_inplace = True
60
- if hasattr(self.model_runner.model, "sampler"):
61
- (self.model_runner.model.sampler.should_modify_greedy_probs_inplace
62
- ) = True
63
-
64
- @torch.inference_mode()
65
- def sampler_output(
66
- self,
67
- execute_model_req: ExecuteModelRequest,
68
- sample_len: int,
69
- seq_ids_with_bonus_token_in_last_step: Set[int],
70
- ) -> Tuple[List[SamplerOutput], bool]:
71
- """Run the model forward pass sample_len times. Returns the list of
72
- sampler output, one per model forward pass, along with indicator of
73
- whether torch tensor in sampler output need to be transposed in latter
74
- sampler_output_to_torch logic.
75
-
76
- For multi step worker, this indicator shall be True.
77
- """
78
- self._raise_if_unsupported(execute_model_req)
79
- # Expand the batch for sequences with a bonus token.
80
- # Perform a forward pass on the expanded batch and filter the
81
- # response to retain only the original sequences' responses.
82
- expanded_request, indices_of_seq_with_bonus_tokens =\
83
- self._expand_execute_model_request(
84
- execute_model_req, seq_ids_with_bonus_token_in_last_step)
85
-
86
- # Run model sample_len times.
87
- model_outputs: List[SamplerOutput] = []
88
- if current_platform.is_cuda_alike() and isinstance(
89
- self.model_runner, TP1DraftModelRunner
90
- ) and self.model_runner.supports_gpu_multi_step(expanded_request):
91
- # Here we run the draft_model_runner with multi-step prepare
92
- # on the GPU directly
93
- expanded_request.num_steps = sample_len
94
- self.model_runner.set_indices_of_seq_with_bonus_tokens(
95
- indices_of_seq_with_bonus_tokens)
96
- model_outputs = self.execute_model(
97
- execute_model_req=expanded_request)
98
- else:
99
- # Here we run multi-step directly, with every step prepared
100
- # on the CPU.
101
- # TODO: Remove this branch once DraftModelRunner supports TP>1
102
- # and other restrictions that are part of DraftModelRunner's
103
- # supports_gpu_multi_step(..)
104
- if expanded_request.previous_hidden_states is not None:
105
- self.worker.model_runner.return_hidden_states = True
106
- for _ in range(sample_len):
107
- model_output: List[SamplerOutput] = self.worker.execute_model(
108
- execute_model_req=expanded_request)
109
- assert (len(model_output) == 1
110
- ), "composing multistep workers not supported"
111
- model_output = model_output[0]
112
- self._maybe_update_previous_hidden_states(
113
- model_output, expanded_request)
114
-
115
- self._append_new_tokens(
116
- model_output, expanded_request.seq_group_metadata_list,
117
- indices_of_seq_with_bonus_tokens)
118
- model_outputs.append(model_output)
119
-
120
- # move indices to device to avoid stream sync
121
- indices_of_seq_with_bonus_tokens = torch.tensor(
122
- indices_of_seq_with_bonus_tokens, device=self.device)
123
- filtered_model_outputs = self._filter_model_output(
124
- model_outputs, indices_of_seq_with_bonus_tokens)
125
- return filtered_model_outputs, True
126
-
127
- @staticmethod
128
- def _maybe_update_previous_hidden_states(
129
- model_output: SamplerOutput,
130
- expanded_request: ExecuteModelRequest) -> None:
131
- """
132
- Updates the previous hidden states in an expanded request
133
- in-place with the hidden states from the model output.
134
- """
135
- if expanded_request.previous_hidden_states is not None:
136
- expanded_request.previous_hidden_states = HiddenStates(
137
- model_output.hidden_states,
138
- expanded_request.seq_group_metadata_list)
139
-
140
- @staticmethod
141
- def _expand_execute_model_request(
142
- execute_model_req: ExecuteModelRequest,
143
- seq_with_bonus_token_in_last_step: set,
144
- ) -> Tuple[ExecuteModelRequest, List[int]]:
145
- """
146
- Expands the execute model request based on sequences with bonus
147
- tokens.
148
-
149
- For each sequence with a bonus token, this method creates a new
150
- sequence without the bonus token and adds it to the execute model
151
- request. The original sequence groups are also retained. The indices
152
- of the original sequence groups are returned for further processing.
153
-
154
- Args:
155
- execute_model_req (ExecuteModelRequest): The original execute
156
- model request.
157
- seq_with_bonus_token_in_last_step (set): Set of sequence IDs that
158
- contain bonus tokens.
159
-
160
- Returns:
161
- Tuple[ExecuteModelRequest, List[int]]: The updated execute model
162
- request with expanded sequences and a list of indices corresponding
163
- to the original sequence groups.
164
- """
165
- updated_seq_group_metadata_list: List[SequenceGroupMetadata] = []
166
- updated_execute_model_req = execute_model_req.clone(
167
- updated_seq_group_metadata_list)
168
- indices_of_original_sequence_groups = []
169
- for seq_group in execute_model_req.seq_group_metadata_list:
170
- seq_group_has_bonus_tokens = False
171
- for seq_id, _ in seq_group.seq_data.items():
172
- # Identify sequences with bonus tokens in the sequence group.
173
- if seq_id in seq_with_bonus_token_in_last_step:
174
- seq_group_has_bonus_tokens = True
175
- break
176
- if seq_group_has_bonus_tokens:
177
- #Create new sequences without the last bonus token. These new
178
- # sequence have the same sequence id as the original sequence.
179
- # We create a new sequence group and add them there.
180
- updated_seq_group_without_bonus_token = \
181
- MultiStepWorker._copy_seq_metadata_excluding_last_token(
182
- seq_group, seq_with_bonus_token_in_last_step)
183
- updated_seq_group_metadata_list.append(
184
- updated_seq_group_without_bonus_token)
185
- # Add the original sequence group.
186
- updated_seq_group_metadata_list.append(
187
- MultiStepWorker._shallow_copy_seq_group_metadata(seq_group))
188
- # Record the index of the original sequence group.
189
- indices_of_original_sequence_groups.append(
190
- len(updated_seq_group_metadata_list) - 1)
191
-
192
- updated_execute_model_req.seq_group_metadata_list =\
193
- updated_seq_group_metadata_list
194
-
195
- if isinstance(updated_execute_model_req.previous_hidden_states,
196
- HiddenStates):
197
- updated_execute_model_req.previous_hidden_states\
198
- .expand_with_bonus_tokens(seq_with_bonus_token_in_last_step)
199
-
200
- return updated_execute_model_req, indices_of_original_sequence_groups
201
-
202
- @staticmethod
203
- def _filter_model_output(
204
- expanded_batch_outputs: List[SamplerOutput],
205
- output_indices_to_retain: torch.Tensor) -> List[SamplerOutput]:
206
- """
207
- Filters the model output to include only the specified sequence
208
- outputs. This method contracts the expanded batch output from the
209
- model to retain the outputs of only those sequences indicated by the
210
- provided indices.
211
-
212
- Args:
213
- expanded_batch_output (List[SamplerOutput]): The expanded output
214
- batch from the model.
215
- output_indices_to_retain (torch.Tensor): Indices of the model
216
- outputs to retain.
217
-
218
- Returns:
219
- List[SamplerOutput]: A list containing the filtered model
220
- outputs for the specified indices.
221
- """
222
- return [
223
- SamplerOutput(
224
- outputs=[
225
- expanded_batch_output.outputs[i]
226
- for i in output_indices_to_retain
227
- ] if len(expanded_batch_output.outputs) > 0 else [],
228
- sampled_token_probs=(
229
- expanded_batch_output.
230
- sampled_token_probs[output_indices_to_retain]
231
- if expanded_batch_output.sampled_token_probs is not None
232
- else None),
233
- logprobs=(
234
- expanded_batch_output.logprobs[output_indices_to_retain]
235
- if expanded_batch_output.logprobs is not None else None),
236
- sampled_token_ids=(expanded_batch_output.
237
- sampled_token_ids[output_indices_to_retain]
238
- if expanded_batch_output.sampled_token_ids
239
- is not None else None))
240
- for expanded_batch_output in expanded_batch_outputs
241
- ]
242
-
243
- def get_spec_proposals(
244
- self,
245
- execute_model_req: ExecuteModelRequest,
246
- seq_ids_with_bonus_token_in_last_step: set,
247
- ) -> SpeculativeProposals:
248
- """Produce speculations given an input batch of sequences. The number of
249
- speculative tokens per sequence is determined by max_proposal_len.
250
- """
251
- return self._proposer.get_spec_proposals(
252
- execute_model_req, seq_ids_with_bonus_token_in_last_step)
253
-
254
- @staticmethod
255
- def _append_new_tokens(
256
- model_output: List[SamplerOutput],
257
- seq_group_metadata_list: List[SequenceGroupMetadata],
258
- indices_of_seq_with_bonus_tokens: List[int]) -> None:
259
- """Given model output from a single run, append the tokens to the
260
- sequences. This is normally done outside of the worker, but it is
261
- required if the worker is to perform multiple forward passes.
262
- """
263
- count = 0
264
- for index, (seq_group_metadata, sequence_group_outputs) in enumerate(
265
- zip(seq_group_metadata_list, model_output)):
266
- seq_group_metadata.is_prompt = False
267
-
268
- for seq_output in sequence_group_outputs.samples:
269
- # NOTE: Beam search is not supported, so we can assume that
270
- # parent_seq_id == seq_id.
271
- seq = seq_group_metadata.seq_data[seq_output.parent_seq_id]
272
-
273
- token_id = seq_output.output_token
274
- token_logprob = seq_output.logprobs[token_id]
275
- # Determine the actual token ID to be generated,
276
- # considering bonus tokens
277
- if index != indices_of_seq_with_bonus_tokens[count]:
278
- bonus_seq_metadata = seq_group_metadata_list[
279
- indices_of_seq_with_bonus_tokens[count]]
280
- _, bonus_token_seq_data = next(
281
- iter(bonus_seq_metadata.seq_data.items()))
282
- token_id = bonus_token_seq_data.output_token_ids[-1]
283
- else:
284
- count += 1
285
-
286
- seq.append_token_id(token_id, token_logprob.logprob,
287
- seq_output.output_embed)
288
- seq.update_num_computed_tokens(1)
289
-
290
- @staticmethod
291
- def _shallow_copy_seq_group_metadata(
292
- seq_group_metadata: SequenceGroupMetadata, ) -> SequenceGroupMetadata:
293
- """Copy input data structures to remove side-effects when input data
294
- structures are shared with other modules.
295
-
296
- Helpful when the vLLM scheduler runs in the same process as the worker.
297
- The alternative is deep-copying (or other form of deep copy); this has
298
- performance downsides.
299
- """
300
- # Shallow-copy the SequenceGroupMetadata. This allows us to
301
- # append tokens and change is_prompt without external side-effects.
302
- # We must shallow-copy seq_group_metadata as is_prompt could change.
303
- new_seq_group_metadata = copy.copy(seq_group_metadata)
304
-
305
- # We must shallow-copy seq_data as we will append token ids
306
- new_seq_data: Dict[int, SequenceData] = {}
307
- for seq_id, old_seq_data in seq_group_metadata.seq_data.items():
308
- new_seq_data[seq_id] = copy.copy(old_seq_data)
309
- new_seq_data[seq_id].output_token_ids =\
310
- old_seq_data.output_token_ids[:]
311
-
312
- new_seq_group_metadata.seq_data = new_seq_data
313
- return new_seq_group_metadata
314
-
315
- @staticmethod
316
- def _copy_seq_metadata_excluding_last_token(
317
- seq_group_metadata: SequenceGroupMetadata,
318
- seq_ids_to_copy: Set[int],
319
- ) -> SequenceGroupMetadata:
320
- """
321
- Creates a shallow copy of the given SequenceGroupMetadata, retaining
322
- only the sequence IDs specified in seq_ids_to_copy. For each of these
323
- sequence IDs, all output_token_ids except the last one are copied.
324
- Sequence IDs not in seq_ids_to_copy are excluded from the copy.
325
-
326
- Parameters:
327
- seq_group_metadata (SequenceGroupMetadata): The original sequence
328
- group metadata.
329
- seq_ids_to_copy (Set[int]): The set of sequence IDs to include in the
330
- copy.
331
-
332
- Returns:
333
- SequenceGroupMetadata: A shallow copy of the sequence group metadata
334
- with the specified modifications.
335
- """
336
- # Shallow-copy the SequenceGroupMetadata.
337
- new_seq_group_metadata = copy.copy(seq_group_metadata)
338
- # Shallow-copy seq_data and modify the output_token_ids.
339
- new_seq_data: Dict[int, SequenceData] = {}
340
- for seq_id, old_seq_data in seq_group_metadata.seq_data.items():
341
- if (seq_id in seq_ids_to_copy):
342
- new_seq_data[seq_id] = copy.copy(old_seq_data)
343
- # Copy all the output token ids except the last.
344
- # Also reduce num_computed_tokens by 1 since we are not
345
- # including the last output token.
346
- # NOTE: num_computed_tokens is not directly used by the
347
- # speculative decoding workers, as it is only relevant for
348
- # chunked prefill, which is disabled for speculative decoding.
349
- # However, to maintain consistency in num_computed_tokens,
350
- # we update it here.
351
- new_seq_data[seq_id].output_token_ids =\
352
- old_seq_data.output_token_ids[:-1]
353
- new_seq_data[seq_id].update_num_computed_tokens(-1)
354
- new_seq_group_metadata.seq_data = new_seq_data
355
- return new_seq_group_metadata
356
-
357
- def _assert_enough_kv_space(
358
- self, seq_group_metadata_list: List[SequenceGroupMetadata],
359
- num_steps: int) -> None:
360
- """Assert there are enough physical blocks per sequence to store the
361
- current KV plus additional KV from num_steps tokens.
362
- """
363
- assert self.model_runner.block_size is not None
364
- for seq_group_metadata in seq_group_metadata_list:
365
- # Only one seq_id is guaranteed because there is no beam search.
366
- seq_id = list(seq_group_metadata.seq_data.keys())[0]
367
- seq = seq_group_metadata.seq_data[seq_id]
368
-
369
- # After num_steps, the seq len will be the current seq len
370
- # plus one token per step.
371
- final_seq_len = seq.get_len() + num_steps
372
-
373
- # We will have final_seq_len - 1 KV because vLLM saves KV for a
374
- # token in the iteration after the token was generated.
375
- required_num_kv_slots = final_seq_len - 1
376
-
377
- # The allocated number of kv slots is the number of allocated blocks
378
- # times the number of slots of block.
379
- number_physical_blocks = len(
380
- seq_group_metadata.block_tables[seq_id])
381
- allocated_kv_slots = (number_physical_blocks *
382
- self.model_runner.block_size)
383
-
384
- if required_num_kv_slots > allocated_kv_slots:
385
- request_id = seq_group_metadata.request_id
386
- raise ValueError(
387
- "The worker attempted to run "
388
- f"{num_steps} times but found insufficient KV space for "
389
- f"{request_id=} {seq_id=}. ({allocated_kv_slots=} "
390
- f"{required_num_kv_slots=}).")
391
-
392
- def _raise_if_unsupported(
393
- self,
394
- execute_model_req: ExecuteModelRequest,
395
- ) -> None:
396
- """MultiStepWorker does not yet implement support for cache swap
397
- operations or beam search.
398
- """
399
- if any([
400
- execute_model_req.blocks_to_swap_in,
401
- execute_model_req.blocks_to_swap_out,
402
- execute_model_req.blocks_to_copy
403
- ]):
404
- raise NotImplementedError(
405
- "MultiStepWorker does not support cache operations")
406
-
407
- if any(
408
- len(seq_group_metadata.seq_data.keys()) != 1
409
- for seq_group_metadata in
410
- execute_model_req.seq_group_metadata_list):
411
- raise NotImplementedError(
412
- "MultiStepWorker does not support beam search.")
413
-
414
- def maybe_load_lm_head_weight(
415
- self,
416
- lm_head_weight: torch.Tensor,
417
- ) -> None:
418
- weight_loader = getattr(
419
- self.worker.model_runner.model_runner.model.lm_head.weight,
420
- "weight_loader", default_weight_loader)
421
- weight_loader(
422
- self.worker.model_runner.model_runner.model.lm_head.weight,
423
- lm_head_weight)
vllm/spec_decode/ngram_worker.py
deleted
-196
@@ -1,196 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import weakref
5
-from typing import List, Optional, Set, Tuple
6
-
7
-import torch
8
-import torch.nn as nn
9
-
10
-from vllm.config import VllmConfig
11
-from vllm.model_executor.layers.sampler import SamplerOutput
12
-from vllm.sequence import ExecuteModelRequest
13
-from vllm.spec_decode.interfaces import SpeculativeProposals
14
-from vllm.spec_decode.proposer_worker_base import NonLLMProposerWorkerBase
15
-from vllm.spec_decode.top1_proposer import Top1Proposer
16
-
17
-
18
-class _DummyModel(nn.Module):
19
- pass
20
-
21
-
22
-class NGramWorker(NonLLMProposerWorkerBase):
23
- """NGramWorker provides a light drafter without need for model.
24
-
25
- Current NGramWorker only implements prompt lookup decoding,
26
- and in future we may also do RAG type drafter and other scenarios
27
- which don't rely on LLM model to give proposals.
28
- """
29
-
30
- def __init__(
31
- self,
32
- vllm_config: VllmConfig,
33
- local_rank: int,
34
- device_type: str = "cuda",
35
- **kwargs,
36
- ):
37
- super().__init__(vllm_config)
38
-
39
- # Get local_rank/vocab_size from kwargs attribute
40
- self.local_rank = local_rank
41
- self.device_type = device_type
42
-
43
- # Lazy initialization list.
44
- self._proposer: Top1Proposer
45
-
46
- def set_ngram_window_size(self, ngram_prompt_lookup_min: int,
47
- ngram_prompt_lookup_max: int):
48
- # Search valid candidate window between
49
- # ngram_prompt_lookup_min/ngram_prompt_lookup_max
50
- self.ngram_prompt_lookup_max = ngram_prompt_lookup_max
51
- self.ngram_prompt_lookup_min = ngram_prompt_lookup_min
52
-
53
- def init_device(self):
54
- self.device = torch.device(f"{self.device_type}:{self.local_rank}")
55
-
56
- # Current NGramWorker only supports Top1Proposer
57
- self._proposer = Top1Proposer(
58
- weakref.proxy(self), # type: ignore[arg-type]
59
- device=self.device,
60
- vocab_size=self.vocab_size,
61
- )
62
-
63
- def load_model(self) -> None:
64
- pass # Dummy
65
-
66
- def get_model(self) -> nn.Module:
67
- return _DummyModel()
68
-
69
- def sampler_output(
70
- self,
71
- execute_model_req: ExecuteModelRequest,
72
- sample_len: int,
73
- # Unused parameter. NGramWorker does not use the KV Cache and
74
- # therefore does not need this parameter.
75
- seq_ids_with_bonus_token_in_last_step: Set[int],
76
- ) -> Tuple[Optional[List[Optional[SamplerOutput]]], bool]:
77
- """NGram match algo to pick proposal candidate. Returns the list of
78
- sampler output, one per SequenceGroupMetadata.
79
-
80
- For ngram worker, we already done needed transposed internal, so the
81
- indicator pass to sampler_output_to_torch shall be False.
82
- """
83
- self._raise_if_unsupported(execute_model_req)
84
-
85
- has_spec_out = False
86
- token_id_list: List[Optional[torch.Tensor]] = []
87
- token_prob_list: List[Optional[torch.Tensor]] = []
88
- for idx, seq_group_metadata in enumerate(
89
- execute_model_req.seq_group_metadata_list):
90
- seq_data = next(iter(seq_group_metadata.seq_data.values()))
91
-
92
- seq_len = seq_data.get_len()
93
- # When seq_len is less than 3072 (3K), we use CPU to perform
94
- # the ngram match. Otherwise, we use the device specified in
95
- # the model config (normally GPU). 3072 is a rough threshold
96
- # based on profiling on H100, and it can be adjusted based
97
- # on the actual performance on different hardware.
98
- cur_device = "cpu" if seq_len < 3072 else self.device
99
- input_ids = torch.as_tensor(seq_data.get_token_ids(),
100
- dtype=torch.long,
101
- device=cur_device)
102
- input_length = seq_data.get_len()
103
-
104
- for ngram_size in range(
105
- min(self.ngram_prompt_lookup_max, input_length - 1),
106
- self.ngram_prompt_lookup_min - 1,
107
- -1,
108
- ):
109
- ngram_tensor = input_ids[-ngram_size:]
110
- if ngram_size == 1:
111
- # Do not match itself and do not use unfold and all
112
- matches = (input_ids[:-1] == ngram_tensor)
113
- else:
114
- windows = input_ids.unfold(dimension=0,
115
- size=ngram_size,
116
- step=1)
117
- # Do not match itself
118
- matches = (windows[:-1] == ngram_tensor).all(dim=-1)
119
-
120
- # first_match includes "values" (bool), indicating whether
121
- # the match is found, and "indices", indicating the index
122
- # of the first match.
123
- first_match = matches.max(dim=-1)
124
- if first_match.values.item():
125
- proposal_start_idx = first_match.indices.add_(ngram_size)
126
- spec_indices = (
127
- proposal_start_idx).repeat(sample_len) + torch.arange(
128
- sample_len, device=cur_device)
129
- spec_indices.clamp_(max=input_ids.shape[-1] - 1)
130
- res = input_ids.gather(dim=-1,
131
- index=spec_indices).to(self.device)
132
- token_id_list.append(res)
133
- token_prob_list.append(
134
- torch.nn.functional.one_hot(
135
- res,
136
- num_classes=self.vocab_size).to(torch.float32))
137
- has_spec_out = True
138
- break
139
- else:
140
- token_id_list.append(None)
141
- token_prob_list.append(None)
142
-
143
- if not has_spec_out:
144
- return None, False
145
-
146
- outputs: List[Optional[SamplerOutput]] = []
147
- for idx in range(len(execute_model_req.seq_group_metadata_list)):
148
- if token_id_list[idx] is None:
149
- outputs.append(None)
150
- else:
151
- outputs.append(
152
- SamplerOutput(
153
- outputs=None,
154
- sampled_token_probs=token_prob_list[idx],
155
- logprobs=torch.zeros((sample_len, self.vocab_size),
156
- dtype=torch.float32,
157
- device=self.device),
158
- sampled_token_ids=token_id_list[idx],
159
- ))
160
-
161
- return outputs, False
162
-
163
- def get_spec_proposals(
164
- self,
165
- execute_model_req: ExecuteModelRequest,
166
- # Unused parameter. NGramWorker does not use the KV Cache and
167
- # therefore does not need this parameter.
168
- seq_ids_with_bonus_token_in_last_step: Set[int],
169
- ) -> SpeculativeProposals:
170
- """Produce speculations given an input batch of sequences. The number of
171
- speculative tokens per sequence is determined by max_proposal_len.
172
- """
173
- return self._proposer.get_spec_proposals(
174
- execute_model_req, seq_ids_with_bonus_token_in_last_step)
175
-
176
- def _raise_if_unsupported(
177
- self,
178
- execute_model_req: ExecuteModelRequest,
179
- ) -> None:
180
- """NGramWorker does not yet implement support for cache swap
181
- operations or beam search.
182
- """
183
- if any([
184
- execute_model_req.blocks_to_swap_in,
185
- execute_model_req.blocks_to_swap_out,
186
- execute_model_req.blocks_to_copy
187
- ]):
188
- raise NotImplementedError(
189
- "NGramWorker does not support cache operations")
190
-
191
- if any(
192
- len(seq_group_metadata.seq_data.keys()) != 1
193
- for seq_group_metadata in
194
- execute_model_req.seq_group_metadata_list):
195
- raise NotImplementedError(
196
- "NGramWorker does not support beam search.")
vllm/spec_decode/proposer_worker_base.py
deleted
-59
@@ -1,59 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from abc import ABC, abstractmethod
5
-from typing import List, Optional, Set, Tuple
6
-
7
-from vllm.model_executor.layers.sampler import SamplerOutput
8
-from vllm.sequence import ExecuteModelRequest
9
-from vllm.spec_decode.interfaces import SpeculativeProposer
10
-from vllm.worker.worker_base import LoRANotSupportedWorkerBase
11
-
12
-
13
-class ProposerWorkerBase(LoRANotSupportedWorkerBase, SpeculativeProposer):
14
- """Interface for proposer workers"""
15
-
16
- @abstractmethod
17
- def sampler_output(
18
- self,
19
- execute_model_req: ExecuteModelRequest,
20
- sample_len: int,
21
- # A set containing all sequence IDs that were assigned bonus tokens
22
- # in their last forward pass. This set is used to backfill the KV cache
23
- # with the key-value pairs of the penultimate token in the sequences.
24
- # This parameter is only used by the MultiStepWorker, which relies on
25
- # the KV cache for token generation. It is not used by workers that
26
- # do not utilize the KV cache.
27
- seq_ids_with_bonus_token_in_last_step: Set[int]
28
- ) -> Tuple[Optional[List[SamplerOutput]], bool]:
29
- raise NotImplementedError
30
-
31
- def set_include_gpu_probs_tensor(self) -> None:
32
- """Implementation optional"""
33
- pass
34
-
35
- def set_should_modify_greedy_probs_inplace(self) -> None:
36
- """Implementation optional"""
37
- pass
38
-
39
-
40
-class NonLLMProposerWorkerBase(ProposerWorkerBase, ABC):
41
- """Proposer worker which does not use a model with kvcache"""
42
-
43
- def execute_model(
44
- self,
45
- execute_model_req: Optional[ExecuteModelRequest] = None
46
- ) -> List[SamplerOutput]:
47
- """get_spec_proposals is used to get the proposals"""
48
- return []
49
-
50
- def determine_num_available_blocks(self) -> Tuple[int, int]:
51
- """This is never called on the proposer, only the target model"""
52
- raise NotImplementedError
53
-
54
- def initialize_cache(self, num_gpu_blocks: int,
55
- num_cpu_blocks: int) -> None:
56
- pass
57
-
58
- def get_cache_block_size_bytes(self) -> int:
59
- return 0
vllm/spec_decode/smaller_tp_proposer_worker.py
deleted
-196
@@ -1,196 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from typing import List, Optional, Set, Tuple
5
-
6
-import torch
7
-import torch.nn as nn
8
-
9
-from vllm.distributed.parallel_state import (get_tp_group,
10
- init_model_parallel_group,
11
- patch_tensor_parallel_group)
12
-from vllm.logger import init_logger
13
-from vllm.model_executor.layers.sampler import SamplerOutput
14
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
15
-from vllm.sequence import ExecuteModelRequest
16
-from vllm.spec_decode.interfaces import SpeculativeProposals
17
-from vllm.spec_decode.multi_step_worker import MultiStepWorker
18
-from vllm.spec_decode.proposer_worker_base import ProposerWorkerBase
19
-
20
-logger = init_logger(__name__)
21
-
22
-
23
-class _DummyModel(nn.Module):
24
- pass
25
-
26
-
27
-class SmallerTpProposerWorker(ProposerWorkerBase):
28
- """Class which allows a speculative draft model to run with smaller tensor
29
- parallel degree than target model.
30
- This reduces the communication overhead of small draft models.
31
-
32
- To implement this feature, this class differs behavior based on is_dummy
33
- flag, where dummy means worker that does not participate draft generation.
34
- Participating workers use a smaller tp group by patching vLLM's tensor
35
- parallel group temporarily during forward passes of draft models.
36
- """
37
-
38
- @classmethod
39
- def maybe_wrap_worker(cls, worker, draft_tensor_parallel_size: int,
40
- target_tensor_parallel_size: int):
41
- """Wrap the worker in a SmallerTpProposerWorker if necessary.
42
- """
43
- if draft_tensor_parallel_size == target_tensor_parallel_size:
44
- return worker
45
-
46
- # gpu ranks that will generate draft tokens together
47
- draft_ranks = list(range(draft_tensor_parallel_size))
48
-
49
- logger.info("Wrapping {%s} in {%s}", type(worker), cls)
50
- return cls(worker, draft_ranks)
51
-
52
- def __init__(self, worker: MultiStepWorker, draft_ranks: List[int]):
53
- """Create a SmallerTpProposerWorker.
54
-
55
- Args:
56
- worker (~vllm.spec_decode.multi_step_worker.MultiStepWorker): an
57
- actual worker wrapped with this class
58
- draft_ranks (List[int]): if this value is given, only the GPU ranks
59
- written in this value participate in draft generation
60
- """
61
- self._worker = worker
62
- self._draft_ranks = draft_ranks
63
-
64
- # init during init_device
65
- self._is_dummy = False
66
- self._tp_group = None
67
-
68
- def _patch_tensor_parallel_group(self):
69
- """Temporarily patch the global tp group state with its own tp group
70
- state.
71
- """
72
- return patch_tensor_parallel_group(self._tp_group)
73
-
74
- def init_device(self) -> None:
75
- self._is_dummy = get_tp_group().rank not in self._draft_ranks
76
-
77
- # dummy workers do nothing
78
- if self._is_dummy:
79
- return
80
-
81
- # creates tp process group containing only a subset of gpu ranks
82
- local_rank = get_tp_group().local_rank
83
- tp_backend = torch.distributed.get_backend(get_tp_group().device_group)
84
- self._tp_group = init_model_parallel_group([self._draft_ranks],
85
- local_rank, tp_backend)
86
-
87
- with self._patch_tensor_parallel_group():
88
- self._worker.init_device()
89
-
90
- def set_include_gpu_probs_tensor(self) -> None:
91
- if self._is_dummy:
92
- return
93
-
94
- # Need include_gpu_probs_tensor for multi_step_worker
95
- self._worker.set_include_gpu_probs_tensor()
96
-
97
- def set_should_modify_greedy_probs_inplace(self) -> None:
98
- if self._is_dummy:
99
- return
100
-
101
- self._worker.set_should_modify_greedy_probs_inplace()
102
-
103
- def load_model(self) -> None:
104
- if self._is_dummy:
105
- return
106
-
107
- with self._patch_tensor_parallel_group():
108
- self._worker.load_model()
109
-
110
- def determine_num_available_blocks(self) -> Tuple[int, int]:
111
- if self._is_dummy:
112
- # this case is not used now
113
- return -1, -1
114
-
115
- with self._patch_tensor_parallel_group():
116
- return self._worker.determine_num_available_blocks()
117
-
118
- def initialize_cache(self, num_gpu_blocks: int,
119
- num_cpu_blocks: int) -> None:
120
- if self._is_dummy:
121
- return
122
-
123
- with self._patch_tensor_parallel_group():
124
- self._worker.initialize_cache(num_gpu_blocks, num_cpu_blocks)
125
-
126
- def sampler_output(
127
- self,
128
- execute_model_req: ExecuteModelRequest,
129
- sample_len: int,
130
- seq_ids_with_bonus_token_in_last_step: Set[int],
131
- ) -> Tuple[List[SamplerOutput], bool]:
132
- # Do not check _is_dummy, as it's always called by get_spec_proposals
133
- return self._worker.sampler_output(
134
- execute_model_req, sample_len,
135
- seq_ids_with_bonus_token_in_last_step)
136
-
137
- def get_spec_proposals(
138
- self,
139
- execute_model_req: ExecuteModelRequest,
140
- seq_ids_with_bonus_token_in_last_step: Set[int],
141
- ) -> SpeculativeProposals:
142
- """Produce speculations given an input batch of sequences. The number of
143
- speculative tokens per sequence is determined by max_proposal_len.
144
- """
145
- if self._is_dummy:
146
- return SpeculativeProposals(None, None, None)
147
-
148
- with self._patch_tensor_parallel_group():
149
- return self._worker.get_spec_proposals(
150
- execute_model_req, seq_ids_with_bonus_token_in_last_step)
151
-
152
- def get_model(self) -> nn.Module:
153
- if self._is_dummy:
154
- return _DummyModel()
155
-
156
- with self._patch_tensor_parallel_group():
157
- return self._worker.get_model()
158
-
159
- def execute_model(
160
- self,
161
- execute_model_req: Optional[ExecuteModelRequest] = None
162
- ) -> List[SamplerOutput]:
163
- if self._is_dummy:
164
- return []
165
-
166
- with self._patch_tensor_parallel_group():
167
- return self._worker.execute_model(execute_model_req)
168
-
169
- def get_cache_block_size_bytes(self) -> int:
170
- if self._is_dummy:
171
- # by returning zero, target worker can use the entire kv cache space
172
- return 0
173
-
174
- return self._worker.get_cache_block_size_bytes()
175
-
176
- @property
177
- def vocab_size(self) -> int:
178
- return self._worker.vocab_size
179
-
180
- def maybe_load_lm_head_weight(
181
- self,
182
- lm_head_weight: torch.Tensor,
183
- ) -> None:
184
- if self._is_dummy:
185
- return
186
-
187
- with self._patch_tensor_parallel_group():
188
- weight_loader = getattr(
189
- self._worker.worker.model_runner.model_runner.model.\
190
- lm_head.weight,
191
- "weight_loader",
192
- default_weight_loader)
193
- weight_loader(
194
- self._worker.worker.model_runner.model_runner.model.\
195
- lm_head.weight,
196
- lm_head_weight)
vllm/spec_decode/spec_decode_worker.py
deleted
-1326
@@ -1,1326 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import copy
5
-from collections import defaultdict
6
-from functools import cached_property
7
-from typing import Any, Dict, List, Optional, Set, Tuple, Type
8
-
9
-import torch
10
-import torch.nn as nn
11
-
12
-from vllm.config import ParallelConfig, SpeculativeConfig, VllmConfig
13
-from vllm.distributed.communication_op import (broadcast_tensor_dict,
14
- get_tp_group,
15
- tensor_model_parallel_gather)
16
-from vllm.distributed.parallel_state import model_parallel_is_initialized
17
-from vllm.logger import init_logger
18
-from vllm.model_executor.layers.rejection_sampler import RejectionSampler
19
-from vllm.model_executor.layers.sampler import SamplerOutput
20
-from vllm.model_executor.layers.spec_decode_base_sampler import (
21
- SpecDecodeBaseSampler, SpecDecodeStochasticBaseSampler)
22
-from vllm.model_executor.layers.typical_acceptance_sampler import (
23
- TypicalAcceptanceSampler)
24
-from vllm.platforms import current_platform
25
-from vllm.sequence import (VLLM_INVALID_TOKEN_ID,
26
- CompletionSequenceGroupOutput, ExecuteModelRequest,
27
- HiddenStates, SequenceGroupMetadata,
28
- get_all_seq_ids_and_request_ids)
29
-from vllm.spec_decode.batch_expansion import BatchExpansionTop1Scorer
30
-
31
-if current_platform.is_cuda_alike():
32
- from vllm.spec_decode.draft_model_runner import TP1DraftModelRunner
33
-
34
-from vllm.spec_decode.interfaces import (SpeculativeProposals,
35
- SpeculativeScorer, SpeculativeScores)
36
-from vllm.spec_decode.medusa_worker import MedusaWorker
37
-from vllm.spec_decode.metrics import AsyncMetricsCollector
38
-from vllm.spec_decode.mlp_speculator_worker import MLPSpeculatorWorker
39
-from vllm.spec_decode.mqa_scorer import MQAScorer
40
-from vllm.spec_decode.multi_step_worker import MultiStepWorker
41
-from vllm.spec_decode.ngram_worker import NGramWorker
42
-from vllm.spec_decode.proposer_worker_base import ProposerWorkerBase
43
-from vllm.spec_decode.smaller_tp_proposer_worker import SmallerTpProposerWorker
44
-from vllm.spec_decode.target_model_runner import TargetModelRunner
45
-from vllm.spec_decode.util import (Timer, create_logprobs_output,
46
- create_sequence_group_output,
47
- get_all_num_logprobs,
48
- get_sampled_token_logprobs, nvtx_range,
49
- split_batch_by_proposal_len)
50
-from vllm.utils import resolve_obj_by_qualname
51
-from vllm.worker.worker_base import LoRANotSupportedWorkerBase, WorkerBase
52
-
53
-logger = init_logger(__name__)
54
-
55
-
56
-def create_spec_worker(*args, **kwargs) -> "SpecDecodeWorker":
57
- """Helper method that is the entrypoint for Executors which use
58
- WorkerWrapper. It constructs a SpecDecodeWorker from the speculative config.
59
- """
60
- vllm_config: VllmConfig = kwargs.get("vllm_config")
61
- speculative_config: SpeculativeConfig = vllm_config.speculative_config
62
- assert speculative_config is not None
63
-
64
- if vllm_config.parallel_config.pipeline_parallel_size > 1:
65
- raise NotImplementedError("Speculative decoding is currently "
66
- "incompatible with pipeline parallelism")
67
-
68
- draft_worker_kwargs = kwargs.copy()
69
-
70
- kwargs["model_runner_cls"] = TargetModelRunner
71
- target_worker_config = copy.deepcopy(vllm_config)
72
- target_worker_config.parallel_config.worker_cls =\
73
- target_worker_config.parallel_config.sd_worker_cls
74
- cls = resolve_obj_by_qualname(
75
- target_worker_config.parallel_config.worker_cls)
76
- target_worker = cls(*args, **kwargs)
77
- # Set the disable_logprobs variable in the TargetModelRunner instance
78
- # as per its value specified in the SpeculativeConfig.
79
- target_worker.model_runner.disable_logprobs =\
80
- speculative_config.disable_logprobs
81
-
82
- draft_worker_config = copy.deepcopy(vllm_config)
83
- draft_worker_config.model_config = speculative_config.draft_model_config
84
- draft_worker_config.quant_config = VllmConfig._get_quantization_config(
85
- draft_worker_config.model_config,
86
- vllm_config.load_config,
87
- )
88
- speculative_config.draft_parallel_config.worker_cls =\
89
- draft_worker_config.parallel_config.sd_worker_cls
90
- draft_worker_config.parallel_config = speculative_config.draft_parallel_config # noqa
91
- # TODO allow draft-model specific load config.
92
-
93
- # Override draft-model specific worker args.
94
- draft_worker_kwargs.update(
95
- vllm_config=draft_worker_config,
96
- ngram_prompt_lookup_max=speculative_config.prompt_lookup_max,
97
- ngram_prompt_lookup_min=speculative_config.prompt_lookup_min,
98
- )
99
-
100
- spec_decode_worker = SpecDecodeWorker.create_worker(
101
- scorer_worker=target_worker,
102
- draft_worker_kwargs=draft_worker_kwargs,
103
- disable_mqa_scorer=speculative_config.disable_mqa_scorer,
104
- disable_by_batch_size=speculative_config.disable_by_batch_size,
105
- draft_token_acceptance_method=speculative_config.acceptance_method,
106
- typical_acceptance_sampler_posterior_threshold=speculative_config.
107
- posterior_threshold,
108
- typical_acceptance_sampler_posterior_alpha=speculative_config.
109
- posterior_alpha,
110
- disable_logprobs=speculative_config.disable_logprobs,
111
- disable_log_stats=speculative_config.disable_log_stats,
112
- num_speculative_tokens=speculative_config.num_speculative_tokens,
113
- )
114
-
115
- return spec_decode_worker
116
-
117
-
118
-# Reminder: Please update docs/features/compatibility_matrix.md
119
-# If the feature combo become valid
120
-class SpecDecodeWorker(LoRANotSupportedWorkerBase):
121
- """Worker which implements speculative decoding.
122
-
123
- Speculative decoding reduces decoding per-token latency by using a proposal
124
- method, such as a small draft model, to speculate ahead of a larger LLM. The
125
- probabilities of the speculative tokens are then determined by the larger
126
- LLM, after which some verification routine determines which (if any) of the
127
- speculative tokens are accepted by the larger LLM.
128
-
129
- See https://github.com/vllm-project/vllm/pull/2188 and
130
- https://github.com/vllm-project/vllm/pull/3103 for more info.
131
-
132
- The current implementation has the following limitations:
133
- * Only draft-model proposal is implemented (contributions for more forms are
134
- welcome!).
135
- * Only top-1 proposal and scoring are implemented. Tree-attention is left as
136
- future work.
137
- * All sequences in a batch must have the same proposal length, or zero. This
138
- can be improved by having per-sequence speculation in the future.
139
- * The scoring forward pass is done without an MQA kernel, which is
140
- suboptimal especially as the batch size, proposal length, and sequence
141
- lengths grow. Contributions to add a MQA scoring are welcome once
142
- correctness tests pass.
143
- More info here https://docs.google.com/document/d/1T-JaS2T1NRfdP51qzqpyakoCXxSXTtORppiwaj5asxA/edit.
144
- """
145
-
146
- @classmethod
147
- def create_worker(
148
- cls,
149
- scorer_worker: WorkerBase,
150
- draft_worker_kwargs: Dict[str, Any],
151
- disable_mqa_scorer: bool,
152
- disable_by_batch_size: Optional[int],
153
- draft_token_acceptance_method: str,
154
- typical_acceptance_sampler_posterior_threshold: float,
155
- typical_acceptance_sampler_posterior_alpha: float,
156
- disable_logprobs: bool,
157
- disable_log_stats: bool,
158
- num_speculative_tokens: int,
159
- ) -> "SpecDecodeWorker":
160
-
161
- allow_zero_draft_token_step = True
162
- enable_lm_head_weight_load = False
163
- num_spec_prefill_steps = 1
164
- ngram_prompt_lookup_max = (
165
- draft_worker_kwargs.pop("ngram_prompt_lookup_max"))
166
- ngram_prompt_lookup_min = (
167
- draft_worker_kwargs.pop("ngram_prompt_lookup_min"))
168
- draft_model_config = draft_worker_kwargs["vllm_config"].model_config
169
- draft_parallel_config: ParallelConfig = draft_worker_kwargs[
170
- 'vllm_config'].parallel_config
171
- if ngram_prompt_lookup_max > 0:
172
- draft_worker_kwargs[
173
- "device_type"] = scorer_worker.device_config.device.type
174
- proposer_worker = NGramWorker(**draft_worker_kwargs)
175
- proposer_worker.set_ngram_window_size(ngram_prompt_lookup_min,
176
- ngram_prompt_lookup_max)
177
- else:
178
- draft_tp = draft_parallel_config.tensor_parallel_size
179
- target_tp = scorer_worker.parallel_config.tensor_parallel_size
180
-
181
- if draft_model_config.hf_config.model_type == "mlp_speculator":
182
- proposer_worker = MLPSpeculatorWorker(**draft_worker_kwargs)
183
- elif draft_model_config.hf_config.model_type == "medusa":
184
- proposer_worker = MedusaWorker(**draft_worker_kwargs)
185
- else:
186
- if draft_tp == 1:
187
- if current_platform.is_cuda_alike():
188
- draft_worker_kwargs[
189
- "model_runner_cls"] = TP1DraftModelRunner
190
- else:
191
- if draft_model_config.hf_config.model_type == "eagle":
192
- raise NotImplementedError(
193
- f"{draft_model_config.hf_config.model_type} "
194
- "does not support TP > 1 yet")
195
-
196
- allow_zero_draft_token_step = False
197
-
198
- # Load lm_head weight for eagle in init_device
199
- if draft_model_config.hf_config.model_type == "eagle":
200
- enable_lm_head_weight_load = True
201
-
202
- proposer_worker = MultiStepWorker(**draft_worker_kwargs)
203
- if draft_model_config.hf_config.model_type == "deepseek_mtp":
204
- num_spec_prefill_steps = \
205
- draft_model_config.hf_config.n_predict
206
-
207
- proposer_worker = SmallerTpProposerWorker.maybe_wrap_worker(
208
- proposer_worker, draft_tp, target_tp)
209
-
210
- logger.info("Configuring SpecDecodeWorker with proposer=%s",
211
- type(proposer_worker))
212
-
213
- spec_decode_sampler: SpecDecodeBaseSampler = None
214
- if draft_token_acceptance_method == "rejection_sampler":
215
- spec_decode_sampler = RejectionSampler()
216
- elif draft_token_acceptance_method == "typical_acceptance_sampler":
217
- spec_decode_sampler = TypicalAcceptanceSampler(
218
- posterior_threshold=\
219
- typical_acceptance_sampler_posterior_threshold,
220
- posterior_alpha=typical_acceptance_sampler_posterior_alpha,
221
- )
222
- logger.info(
223
- "[Speculative Decoding] Configuring"
224
- " SpecDecodeWorker with sampler=%s", type(spec_decode_sampler))
225
-
226
- if not disable_mqa_scorer:
227
- if scorer_worker.model_runner.attn_backend.get_name(
228
- ) != "FLASH_ATTN":
229
- disable_mqa_scorer = True
230
- logger.info(
231
- "[Speculative Decoding] Disabling MQA scorer as the "
232
- "MQA is only available with flash attn backend.")
233
-
234
- if draft_model_config and \
235
- draft_model_config.max_model_len < \
236
- scorer_worker.model_config.max_model_len:
237
- disable_mqa_scorer = True
238
- logger.info(
239
- "[Speculative Decoding] Disabling MQA scorer as the "
240
- "draft model max_model_len is smaller than the target "
241
- "model max_model_len.")
242
-
243
- if not scorer_worker.model_runner.model_config.enforce_eager:
244
- disable_mqa_scorer = True
245
- logger.info(
246
- "[Speculative Decoding] Disabling MQA scorer as the "
247
- "target model is not running in eager mode.")
248
-
249
- return SpecDecodeWorker(
250
- proposer_worker,
251
- scorer_worker,
252
- disable_mqa_scorer=disable_mqa_scorer,
253
- disable_logprobs=disable_logprobs,
254
- disable_log_stats=disable_log_stats,
255
- disable_by_batch_size=disable_by_batch_size,
256
- spec_decode_sampler=spec_decode_sampler,
257
- allow_zero_draft_token_step=allow_zero_draft_token_step,
258
- enable_lm_head_weight_load=enable_lm_head_weight_load,
259
- num_spec_prefill_steps=num_spec_prefill_steps)
260
-
261
- def __init__(
262
- self,
263
- proposer_worker: ProposerWorkerBase,
264
- scorer_worker: WorkerBase,
265
- spec_decode_sampler: SpecDecodeBaseSampler,
266
- disable_mqa_scorer: bool = False,
267
- disable_logprobs: bool = False,
268
- disable_log_stats: bool = False,
269
- metrics_collector: Optional[AsyncMetricsCollector] = None,
270
- disable_by_batch_size: Optional[int] = None,
271
- allow_zero_draft_token_step: Optional[bool] = True,
272
- enable_lm_head_weight_load: Optional[bool] = False,
273
- num_spec_prefill_steps: int = 1,
274
- ):
275
- """
276
- Create a SpecDecodeWorker.
277
-
278
- Args:
279
- proposer_worker: A worker that can produce speculative tokens for
280
- sequences.
281
- scorer_worker: A worker that produces probabilities of speculative
282
- tokens according to some base model. Typically a vanilla vLLM
283
- Worker.
284
- spec_decode_sampler: A Torch module used to perform acceptance
285
- sampling of the draft tokens in the verification step of
286
- speculative decoding. Currently we support two different
287
- types of sampler namely RejectionSampler and
288
- TypicalAcceptanceSampler. 'spec_decode_sampler' is either an
289
- instance of RejectionSampler or TypicalAcceptanceSampler.
290
- disable_mqa_scorer: If set to True, disable the MQA scorer and use
291
- the BatchExpansionTop1Scorer instead.
292
- disable_logprobs: If set to True, token log probabilities will
293
- not be output in both the draft worker and the target worker.
294
- If set to False, log probabilities will be output by both.
295
- disable_log_stats: If set to True, disable periodic printing of
296
- speculative stage times.
297
- disable_by_batch_size: If the batch size is larger than this,
298
- disable speculative decoding for new incoming requests.
299
- metrics_collector: Helper class for collecting metrics; can be set
300
- for testing purposes.
301
- allow_zero_draft_token_step: whether to allow a step where the draft
302
- model generates no draft token; should disallow when the tp of
303
- draft model is larger than 1 (TODO: #5814)
304
- enable_lm_head_weight_load: whether to load lm_head weight for
305
- draft models like eagle.
306
- num_spec_prefill_steps: number of speculative prefill steps to run
307
- before the speculative decoding starts. This is only used when
308
- the draft model is a deepseek_mtp model that requires prefill
309
- kv cache separately for each MTP layer.
310
- """
311
- self.proposer_worker = proposer_worker
312
- self.scorer_worker = scorer_worker
313
- scorer_runner = getattr(self.scorer_worker, "model_runner", None)
314
- self.generators = scorer_runner.get_generators(
315
- ) if scorer_runner else None
316
- self.disable_by_batch_size = disable_by_batch_size or float("inf")
317
- self.spec_decode_sampler = spec_decode_sampler
318
- self._allow_zero_draft_token_step = allow_zero_draft_token_step
319
- self._enable_lm_head_weight_load = enable_lm_head_weight_load
320
- self._metrics = AsyncMetricsCollector(
321
- self.spec_decode_sampler
322
- ) if metrics_collector is None else metrics_collector
323
- # Tracks the sequence IDs that received a bonus token ID in
324
- # their last forward pass. Needed only if KV cache is being
325
- # used for token generation such as in the case of MultiStepWorker.
326
- self._seq_with_bonus_token_in_last_step: Set[int] = set()
327
- # Tracks the currently active request ids and the sequence IDs
328
- # corresponding to them
329
- self._request_id_seq_id_mapping: Dict[str, Set[int]] = defaultdict(set)
330
- # Tracks if the proposer worker uses the KV cache or not.
331
-
332
- self.probs_dtype = self.spec_decode_sampler.probs_dtype
333
- self.token_id_dtype = self.spec_decode_sampler.token_id_dtype
334
- # Lazy initialization.
335
- self.scorer: SpeculativeScorer
336
- self.disable_mqa_scorer = disable_mqa_scorer
337
-
338
- # Hidden states from target model to pass to proposer
339
- # in the subsequent step.
340
- self.previous_hidden_states: Optional[HiddenStates] = None
341
- self._disable_logprobs = disable_logprobs
342
- self._disable_log_stats = disable_log_stats
343
- self._num_spec_prefill_steps = num_spec_prefill_steps
344
-
345
- def init_device(self) -> None:
346
- """Initialize both scorer and proposer models.
347
- """
348
- # The scorer worker model is initialized first in case the proposer
349
- # model has a smaller TP degree than the target worker.
350
- self.scorer_worker.init_device()
351
- self.proposer_worker.init_device()
352
-
353
- # NOTE(cade): load_model is not part of the WorkerBase interface.
354
- self.scorer_worker.load_model()
355
- self.proposer_worker.load_model()
356
-
357
- if self._enable_lm_head_weight_load:
358
- # NOTE(Shangming): gather lm_head weight when tp enabled
359
- target_lm_head_weight: torch.Tensor = tensor_model_parallel_gather(
360
- self.scorer_worker.model_runner.model_runner.model.lm_head.\
361
- weight.data,
362
- dim=0,
363
- )
364
-
365
- self.proposer_worker.maybe_load_lm_head_weight(
366
- target_lm_head_weight)
367
-
368
- self._metrics.init_tensors(self.rank, device_type=self.device)
369
- if model_parallel_is_initialized():
370
- self.spec_decode_sampler.init_tensors(get_tp_group().local_rank,
371
- device_type=self.device)
372
- else:
373
- self.spec_decode_sampler.init_tensors(self.rank,
374
- device_type=self.device)
375
-
376
- scorer_cls: Type[SpeculativeScorer]
377
- if self.disable_mqa_scorer:
378
- scorer_cls = BatchExpansionTop1Scorer
379
- logger.info("[Speculative Decoding] Use batch "
380
- "expansion for scoring proposals.")
381
- else:
382
- scorer_cls = MQAScorer
383
- logger.info(
384
- "[Speculative Decoding] Use MQA scorer for scoring proposals.")
385
-
386
- self.scorer = scorer_cls(scorer_worker=self.scorer_worker,
387
- device=self.device,
388
- vocab_size=self._vocab_size)
389
-
390
- self._configure_model_sampler_for_spec_decode()
391
-
392
- def load_model(self, *args, **kwargs):
393
- pass
394
-
395
- def _configure_model_sampler_for_spec_decode(self):
396
- """Configure model sampler to emit GPU tensors. This allows spec decode
397
- to keep data on device without transferring to CPU and serializing,
398
- which significantly reduces overhead of sampling during verification.
399
-
400
- NOTE(cade): This breaks abstraction boundaries pretty badly. The better
401
- design is to have the "move to CPU and serialize" sampling decision be
402
- done outside of the model/sampler; this way the "last-mile" worker
403
- object which interfaces with the scheduler can serialize and incur the
404
- performance hit as necessary. This allows us to run the worker several
405
- iterations in a row without incurring the "move to CPU and serialize"
406
- performance penalty.
407
-
408
- Since this requires a large change to vLLM, we defer it to later and
409
- temporarily accept this broken abstraction boundary.
410
-
411
- NOTE(cade): This will require a special check if the proposer worker
412
- does not have a sampler (e.g. ngram speculation).
413
- """
414
- (self.scorer_worker.model_runner.sampler.include_gpu_probs_tensor
415
- ) = True
416
- (self.scorer_worker.model_runner.sampler.
417
- should_modify_greedy_probs_inplace) = True
418
- self.proposer_worker.set_include_gpu_probs_tensor()
419
- self.proposer_worker.set_should_modify_greedy_probs_inplace()
420
-
421
- def determine_num_available_blocks(self) -> Tuple[int, int]:
422
- """Determine the number of cache blocks to use.
423
-
424
- This is done by profiling the scorer model (which is typically the
425
- larger of the two). Then the total memory which would be used by the
426
- scorer cache is divided evenly between the proposer and scorer model KV,
427
- such that the number of blocks is equal in both KV caches.
428
- """
429
- num_gpu_blocks, num_cpu_blocks = (
430
- self.scorer_worker.determine_num_available_blocks())
431
-
432
- scorer_cache_block_size_bytes = (
433
- self.scorer_worker.get_cache_block_size_bytes())
434
- proposer_cache_block_size_bytes = (
435
- self.proposer_worker.get_cache_block_size_bytes())
436
-
437
- new_num_gpu_blocks = split_num_cache_blocks_evenly(
438
- scorer_cache_block_size_bytes, proposer_cache_block_size_bytes,
439
- num_gpu_blocks)
440
- return new_num_gpu_blocks, num_cpu_blocks
441
-
442
- def initialize_cache(self, num_gpu_blocks: int,
443
- num_cpu_blocks: int) -> None:
444
- """Initialize the cache engine of the scorer and proposer workers.
445
- """
446
- self.scorer_worker.initialize_cache(num_gpu_blocks=num_gpu_blocks,
447
- num_cpu_blocks=num_cpu_blocks)
448
- self.proposer_worker.initialize_cache(num_gpu_blocks=num_gpu_blocks,
449
- num_cpu_blocks=num_cpu_blocks)
450
-
451
- def get_model(self) -> nn.Module:
452
- return self.scorer_worker.get_model()
453
-
454
- @torch.inference_mode()
455
- def execute_model(
456
- self,
457
- execute_model_req: Optional[ExecuteModelRequest] = None
458
- ) -> List[SamplerOutput]:
459
- """Perform speculative decoding on the input batch.
460
- """
461
- if self.rank != self._driver_rank:
462
- self._run_non_driver_rank()
463
- return []
464
-
465
- if execute_model_req is None:
466
- # This signals that there's no more requests to process for now.
467
- # All workers are running infinite loop with broadcast_tensor_dict,
468
- # and it stops the loop when the driver broadcasts an empty input.
469
- # Send an empty input to notify all other workers to stop their
470
- # execution loop.
471
- broadcast_tensor_dict({}, src=0)
472
- return []
473
-
474
- self._track_finished_requests(execute_model_req)
475
- disable_all_speculation = self._should_disable_all_speculation(
476
- execute_model_req)
477
- num_lookahead_slots = execute_model_req.num_lookahead_slots
478
- all_prompt = True
479
- atleast_one_prompt = False
480
- all_zero_spec_tokens = True
481
- for sgm in execute_model_req.seq_group_metadata_list:
482
- all_prompt = all_prompt and sgm.is_prompt
483
- atleast_one_prompt = atleast_one_prompt or sgm.is_prompt
484
- all_zero_spec_tokens = all_zero_spec_tokens and (
485
- sgm.num_speculative_tokens == 0)
486
-
487
- if all_prompt and execute_model_req.seq_group_metadata_list:
488
- assert num_lookahead_slots == 0, (
489
- "Prompt only runs should have num_lookahead_slots equal to 0. "
490
- "This should never happen, please file a bug at "
491
- "https://github.com/vllm-project/vllm/issues")
492
- # Speculative decoding is disabled in the following cases:
493
- # 1. Prefill phase: Speculative decoding is not
494
- # used during the prefill phase.
495
- # 2. Auto-disable enabled: The running queue size exceeds
496
- # the specified threshold.
497
- # 3. No request: There are no requests in the batch, or
498
- # none of the requests in the batch have spec decoding enabled.
499
- # In any of these cases, the proposer and scorer workers
500
- # are called normally.
501
- # We expect `num_speculative_tokens` to be None for prefills.
502
- no_spec = (num_lookahead_slots == 0 or disable_all_speculation
503
- or all_zero_spec_tokens)
504
-
505
- # Broadcast how many lookahead slots are scheduled for this step, and
506
- # whether all speculation is disabled, to all non-driver workers.
507
-
508
- # This is required as if the number of draft model runs changes
509
- # dynamically, the non-driver workers won't know unless we perform a
510
- # communication to inform them.
511
-
512
- # no_spec is used to signal non-driver worker about prefill vs decode
513
- # stage. This is needed to ensure that order of execution of proposer
514
- # and scorer is same in both driver and non-driver workers (i.e.,
515
- # scorer -> proposer for prefill and proposer -> scorer in decode). This
516
- # order is needed to support models like EAGLE that take scorer states
517
- # as inputs.
518
- broadcast_dict = dict(
519
- num_lookahead_slots=num_lookahead_slots,
520
- no_spec=no_spec,
521
- disable_all_speculation=disable_all_speculation,
522
- # When both chunked prefill and speculative decoding are enabled
523
- # it is possible that the same batch contains both prefill
524
- # and decodes. If that happens in the scorer we run the batch
525
- # as one single forward pass. However, in the proposer we
526
- # run them as 2 different batches - one for prefill and
527
- # the other for decodes. The variable indicates to the non-driver
528
- # worker that there are prefills as part of the speculative batch
529
- # and hence it needs to run an extra prefill forward pass.
530
- run_spec_proposer_for_prefill=atleast_one_prompt,
531
- )
532
- broadcast_tensor_dict(broadcast_dict, src=self._driver_rank)
533
-
534
- assert execute_model_req.seq_group_metadata_list is not None, (
535
- "speculative decoding requires non-None seq_group_metadata_list")
536
-
537
- self._maybe_disable_speculative_tokens(
538
- disable_all_speculation, execute_model_req.seq_group_metadata_list)
539
-
540
- if no_spec:
541
- return self._run_no_spec(execute_model_req,
542
- skip_proposer=disable_all_speculation)
543
- return self._run_speculative_decoding_step(execute_model_req,
544
- num_lookahead_slots)
545
-
546
- @torch.inference_mode()
547
- def start_worker_execution_loop(self) -> None:
548
- """Execute model loop to perform speculative decoding
549
- in parallel worker."""
550
- while self._run_non_driver_rank():
551
- pass
552
-
553
- def _should_disable_all_speculation(
554
- self, execute_model_req: ExecuteModelRequest) -> bool:
555
- # When the batch size is too large, disable speculative decoding
556
- # to stop trading off throughput for latency.
557
- return (execute_model_req.running_queue_size
558
- >= self.disable_by_batch_size)
559
-
560
- def _maybe_disable_speculative_tokens(
561
- self, disable_all_speculation: bool,
562
- seq_group_metadata_list: List[SequenceGroupMetadata]) -> None:
563
- if not disable_all_speculation:
564
- return
565
-
566
- for seq_group_metadata in seq_group_metadata_list:
567
- # Once num_speculative_tokens is set to 0, the spec decode
568
- # of this request will be disabled forever.
569
- # TODO(comaniac): We currently store spec decoding specific
570
- # state in the global data structure, but we should maintain
571
- # this state within spec decode worker.
572
- seq_group_metadata.num_speculative_tokens = 0
573
-
574
- def _serialize_sampler_output_no_logprobs(
575
- self, execute_model_req: ExecuteModelRequest,
576
- sampler_output: SamplerOutput) -> List[SamplerOutput]:
577
- """
578
- Creates and returns a `SamplerOutput` with only the token IDs being
579
- serialized to CPU and populated in `CompletionSequenceGroupOutput`.
580
- All other parameters in `CompletionSequenceGroupOutput` related to log
581
- probabilities are skipped.
582
-
583
- Args:
584
- execute_model_req (ExecuteModelRequest): The model request that
585
- was executed.
586
- sampler_output (SamplerOutput): The output from the sampler with
587
- only GPU tensors populated.
588
-
589
- Returns:
590
- SamplerOutput: A new `SamplerOutput` instance containing a list of
591
- `CompletionSequenceGroupOutput` objects with only token IDs
592
- populated.
593
- """
594
- seq_output_prompt_logprobs = [
595
- seq.is_prompt and seq.sampling_params.prompt_logprobs is not None
596
- and seq.sampling_params.prompt_logprobs > 0
597
- for seq in execute_model_req.seq_group_metadata_list
598
- ]
599
- # ignore slots for prompt tokens that are filled with INVALID_TOKEN_ID
600
- sampled_token_ids_list = (sampler_output.sampled_token_ids[torch.where(
601
- # subtracting is faster than testing for equality
602
- sampler_output.sampled_token_ids - VLLM_INVALID_TOKEN_ID)[0]] \
603
- if any(seq_output_prompt_logprobs) else \
604
- sampler_output.sampled_token_ids).tolist()
605
-
606
- seq_data_entries = [
607
- (seq_id, seq_data) for sg in \
608
- execute_model_req.seq_group_metadata_list \
609
- for seq_id, seq_data in sg.seq_data.items()
610
- ]
611
- completion_seq_group_output_list: List[
612
- CompletionSequenceGroupOutput] = []
613
- output_index = 0
614
- # Make sure the non-terminal prefill chunks are still aligned with
615
- # their own empty output.
616
- for idx, seq_group_meta in enumerate(
617
- execute_model_req.seq_group_metadata_list):
618
- needs_prompt_logprobs = seq_output_prompt_logprobs[idx]
619
- seq_id, seq_data = seq_data_entries[idx]
620
- if needs_prompt_logprobs:
621
- prompt_token_ids = seq_data.get_prompt_token_ids()
622
-
623
- # Some of these sequences may belong to non-terminal chunks,
624
- # which may still have to report logprobs for prompts.
625
- start = 1 if seq_data._num_computed_tokens == 0 \
626
- else seq_data._num_computed_tokens
627
- end = (seq_data._num_computed_tokens + \
628
- seq_group_meta.token_chunk_size)
629
- prompt_token_ids = prompt_token_ids[start:end]
630
- prompt_logprobs = [
631
- create_logprobs_output(
632
- token_id=p_token_id,
633
- token_id_logprob_rank=-1,
634
- token_id_logprob=0.0,
635
- topk_token_ids=[],
636
- topk_logprobs=[],
637
- ) for p_token_id in prompt_token_ids
638
- ]
639
- else:
640
- prompt_logprobs = None
641
-
642
- # Since we can get chunks here, we dont always have a sampled token
643
- # (only on last chunk) but we still have to provide an output.
644
- if not seq_group_meta.do_sample:
645
- completion_seq_group_output_list.append(
646
- CompletionSequenceGroupOutput(
647
- samples=[], prompt_logprobs=prompt_logprobs))
648
- continue
649
-
650
- # Sequence with output.
651
- completion_seq_group_output_list.append(
652
- create_sequence_group_output(
653
- token_id=sampled_token_ids_list[output_index][0],
654
- token_id_logprob_rank=-1,
655
- token_id_logprob=0.0,
656
- seq_id=seq_id,
657
- topk_token_ids=[],
658
- topk_logprobs=[],
659
- prompt_logprobs=prompt_logprobs))
660
- output_index += 1
661
-
662
- return [SamplerOutput(outputs=completion_seq_group_output_list)]
663
-
664
- @nvtx_range("spec_decode_worker._run_no_spec")
665
- def _run_no_spec(self, execute_model_req: ExecuteModelRequest,
666
- skip_proposer: bool) -> List[SamplerOutput]:
667
- """Run a single generation step without any speculation. The input is
668
- sent to the proposer and scorer model so that the KV cache is consistent
669
- between the two. When skip_proposer is True, the proposer model is
670
- not called, meaning that the kv-cache in proposer for requests is not
671
- updated, so they cannot enable spec decode in the rest decoding.
672
- """
673
-
674
- sampler_output = self.scorer_worker.execute_model(execute_model_req)
675
- assert len(sampler_output) == 1
676
- sampler_output = sampler_output[0]
677
-
678
- # Store hidden states from target model execution, BxD.
679
- hidden_states = sampler_output.hidden_states
680
- if hidden_states is not None:
681
- # Only decodes and prefill terminal chunks need a hidden state.
682
- seq_group_meta_with_hidden = [
683
- sg for sg in execute_model_req.seq_group_metadata_list
684
- if sg.do_sample
685
- ]
686
- if any(seq.is_prompt for seq in seq_group_meta_with_hidden):
687
- # Drop hidden_states with no prediction (eg non-terminal chunks)
688
- hidden_states = hidden_states[
689
- torch.where(sampler_output.sampled_token_ids -
690
- VLLM_INVALID_TOKEN_ID)[0]]
691
- if self.previous_hidden_states is None and len(
692
- seq_group_meta_with_hidden):
693
- self.previous_hidden_states = HiddenStates(
694
- hidden_states, seq_group_meta_with_hidden)
695
- elif self.previous_hidden_states and len(
696
- seq_group_meta_with_hidden):
697
- self.previous_hidden_states.update(hidden_states,
698
- seq_group_meta_with_hidden)
699
- self.previous_hidden_states.prune(seq_group_meta_with_hidden)
700
-
701
- if not skip_proposer:
702
- # We prepare the prefill hidden states here so that there no
703
- # additional complexity in worker for spec_decode vs non_spec_decode
704
- # flow and execute_model doesn't need additional modifications.
705
- execute_model_req.previous_hidden_states = \
706
- prepare_prefill_hidden_states(
707
- sampler_output.prefill_hidden_states)
708
- for i in range(self._num_spec_prefill_steps):
709
- execute_model_req.spec_step_idx = i
710
- self.proposer_worker.execute_model(execute_model_req)
711
-
712
- sampler_output_to_return = (self._serialize_sampler_output_no_logprobs(
713
- execute_model_req=execute_model_req, sampler_output=sampler_output)
714
- if self._disable_logprobs else
715
- [sampler_output])
716
-
717
- # Clear device tensors from sampler output. This reduces communication
718
- # overhead when the engine runs in a different process than the workers.
719
- sampler_output.sampled_token_probs = None
720
- sampler_output.sampled_token_ids = None
721
- sampler_output.logprobs = None
722
- return sampler_output_to_return
723
-
724
- def _run_non_driver_rank(self) -> bool:
725
- """Run proposer and verifier model in non-driver workers. This is used
726
- for both speculation cases (num_lookahead_slots>0) and non-speculation
727
- cases (e.g. prefill).
728
-
729
- Returns True if there are remaining sequences to process.
730
- """
731
- assert self.rank != self._driver_rank
732
-
733
- data = broadcast_tensor_dict(src=self._driver_rank)
734
- if not data:
735
- return False
736
- num_lookahead_slots = data["num_lookahead_slots"]
737
-
738
- # In case of prefill, scorer_worker has to be run before proposer so
739
- # that the hidden states can be propagated to proposer when needed.
740
- if data["no_spec"]:
741
- self.scorer_worker.execute_model()
742
-
743
- if not data["disable_all_speculation"]:
744
- # Even if num_lookahead_slots is zero, we want to run the
745
- # proposer model as it may have KV.
746
- #
747
- # We run the proposer once per lookahead slot. In the future we
748
- # should delegate how many times it runs to the proposer.
749
- for _ in range(max(num_lookahead_slots, 1)):
750
- self.proposer_worker.execute_model()
751
-
752
- if not data["no_spec"]:
753
- self.scorer_worker.execute_model()
754
- if data["run_spec_proposer_for_prefill"]:
755
- self.proposer_worker.execute_model()
756
-
757
- return True
758
-
759
- @nvtx_range("spec_decode_worker._run_speculative_decoding_step")
760
- def _run_speculative_decoding_step(
761
- self, execute_model_req: ExecuteModelRequest,
762
- num_lookahead_slots: int) -> List[SamplerOutput]:
763
- """Execute a single step of speculative decoding.
764
-
765
- This invokes the proposer worker to get k speculative tokens for each
766
- sequence, then scores each speculative token using the scoring worker.
767
-
768
- When `enable_chunked_prefill` is set, scorer will batch decodes and
769
- prefills, while proposer will sync its KV-cache by running an extra
770
- forward on prefills.
771
-
772
- Returns a list of SamplerOutput, each containing a single token per
773
- sequence.
774
- """
775
- # With prefill chunking, expect requests to have prompts first
776
- # so that backend gets prefill|decode.
777
- assert num_lookahead_slots == execute_model_req.num_lookahead_slots
778
-
779
- # Pass last hidden states from target model to proposer
780
- execute_model_req.previous_hidden_states = self.previous_hidden_states
781
- self.previous_hidden_states = None
782
-
783
- with Timer() as proposal_timer:
784
- # Generate proposals using draft worker.
785
- proposals = self.proposer_worker.get_spec_proposals(
786
- execute_model_req, self._seq_with_bonus_token_in_last_step)
787
-
788
- if not self._allow_zero_draft_token_step and proposals.no_proposals:
789
- #TODO: Fix it #5814
790
- raise RuntimeError("Cannot handle cases where distributed draft "
791
- "workers generate no tokens")
792
-
793
- execute_model_req.previous_hidden_states = None
794
-
795
- with Timer() as scoring_timer:
796
- proposal_scores = self.scorer.score_proposals(
797
- execute_model_req,
798
- proposals,
799
- )
800
-
801
- _, (non_spec_seqs, non_spec_indices) = split_batch_by_proposal_len(
802
- execute_model_req.seq_group_metadata_list, proposals.proposal_lens)
803
- # With prefill chunking enabled, `non_spec_seqs` contains prefills too:
804
- # discard decodes that have already been processed by proposer.
805
- non_spec_indices = [
806
- idx for idx in non_spec_indices
807
- if execute_model_req.seq_group_metadata_list[idx].is_prompt
808
- ]
809
- if len(non_spec_indices):
810
- all_hidden_states = proposal_scores.hidden_states
811
- if all_hidden_states is not None:
812
- prefill_hidden_states = all_hidden_states[non_spec_indices]
813
- execute_model_req.previous_hidden_states = \
814
- prepare_prefill_hidden_states(prefill_hidden_states)
815
- # Sync proposer KV cache for prefills.
816
- prefill_req = execute_model_req.clone(non_spec_seqs)
817
- # TODO avoid sampling here?
818
- self.proposer_worker.execute_model(prefill_req)
819
-
820
- with Timer() as verification_timer:
821
- accepted_token_ids, target_logprobs = self._verify_tokens(
822
- execute_model_req.seq_group_metadata_list, proposal_scores,
823
- proposals, execute_model_req.num_lookahead_slots)
824
-
825
- stage_times = (proposal_timer.elapsed_time_ms / num_lookahead_slots,
826
- scoring_timer.elapsed_time_ms,
827
- verification_timer.elapsed_time_ms)
828
-
829
- return self._create_output_sampler_list(
830
- execute_model_req.seq_group_metadata_list,
831
- accepted_token_ids,
832
- target_logprobs=target_logprobs,
833
- prompt_logprobs=proposal_scores.prompt_logprobs
834
- if not self._disable_logprobs else None,
835
- k=execute_model_req.num_lookahead_slots,
836
- stage_times=stage_times)
837
-
838
- @nvtx_range("spec_decode_worker._verify_tokens")
839
- def _verify_tokens(
840
- self,
841
- seq_group_metadata_list: List[SequenceGroupMetadata],
842
- proposal_scores: SpeculativeScores,
843
- proposals: SpeculativeProposals,
844
- max_proposal_len: int,
845
- ) -> Tuple[torch.Tensor, torch.Tensor]:
846
- """Determine which speculative tokens are accepted using the
847
- probabilities of each token according to the proposer and scorer models.
848
-
849
- Returns a tuple of Tensors, one for the accepted token ids and one for
850
- the logprobs according to the scoring model.
851
- """
852
- proposal_lens_list = proposals.proposal_lens.tolist()
853
-
854
- # vLLM currently only supports proposal lens equal to zero or the batch
855
- # proposal len. This adds some complexity (splitting the batch into spec
856
- # and non spec sequences) and should be removed in the future. It can be
857
- # done by supporting per-sequence proposal lens.
858
- (_, spec_indices), (_, non_spec_indices) = split_batch_by_proposal_len(
859
- seq_group_metadata_list, proposal_lens_list)
860
- original_indices = spec_indices + non_spec_indices
861
-
862
- # Get probabilities of target model, including bonus tokens.
863
- proposal_verifier_probs = proposal_scores.probs[spec_indices]
864
-
865
- # Get non-speculative sampled tokens from target model.
866
- non_spec_token_ids = proposal_scores.token_ids[non_spec_indices]
867
-
868
- # Get bonus tokens from target model.
869
- bonus_token_ids = proposal_scores.token_ids[spec_indices, -1:]
870
-
871
- # Get probabilities according to proposal method.
872
- proposal_probs = proposals.proposal_probs[spec_indices]
873
-
874
- # Get proposed tokens.
875
- proposal_token_ids = proposals.proposal_token_ids[spec_indices]
876
-
877
- # Sampler arguments
878
- sampler_extra_kwargs: Dict[str, Any] = {}
879
- if self.generators and isinstance(self.spec_decode_sampler,
880
- SpecDecodeStochasticBaseSampler):
881
- sampler_extra_kwargs["seeded_seqs"] = {
882
- idx: self.generators[sgm.request_id]
883
- for idx, sgm in enumerate(seq_group_metadata_list)
884
- if sgm.sampling_params.seed is not None
885
- }
886
-
887
- accepted_token_ids = self.spec_decode_sampler(
888
- target_with_bonus_probs=proposal_verifier_probs,
889
- bonus_token_ids=bonus_token_ids,
890
- draft_probs=proposal_probs,
891
- draft_token_ids=proposal_token_ids,
892
- **sampler_extra_kwargs,
893
- )
894
- # Append output tokens from non-speculative sequences to
895
- # the accepted token ids tensor.
896
- non_spec_token_ids = non_spec_token_ids.expand(-1, max_proposal_len +
897
- 1).clone()
898
- non_spec_token_ids[:, 1:] = -1
899
- accepted_token_ids = torch.cat(
900
- [accepted_token_ids, non_spec_token_ids])
901
- logprobs = proposal_scores.logprobs
902
- # Rearrange so that results are in the order of the original seq group
903
- # metadata.
904
- accepted_token_ids[original_indices] = accepted_token_ids.clone()
905
-
906
- # B x K+1 x D
907
- hidden_states = proposal_scores.hidden_states
908
- if hidden_states is not None:
909
- # Only get terminal hidden states for next step
910
- terminal_metadata = [
911
- sg for sg in seq_group_metadata_list if sg.do_sample
912
- ]
913
-
914
- # Contract hidden states based on accepted tokens
915
- hs_size = hidden_states.shape[-1]
916
- accepted_index = accepted_token_ids + 1 # Convert -1 to 0
917
- accepted_index = accepted_index.count_nonzero(dim=1).add_(-1) # b
918
- # Drop non-terminal prefill chunks hidden states.
919
- hidden_states = hidden_states[accepted_index !=
920
- VLLM_INVALID_TOKEN_ID]
921
- accepted_index = accepted_index[accepted_index !=
922
- VLLM_INVALID_TOKEN_ID]
923
- assert len(accepted_index) == hidden_states.shape[0] == len(
924
- terminal_metadata)
925
- index = accepted_index[:, None, None].expand(-1, 1,
926
- hs_size) # b x 1 x d
927
- second_last_token_hidden_states = hidden_states[:, -2] # b x d
928
- hidden_states = hidden_states.gather(1, index).squeeze(1) # b x d
929
- # Store hidden states from target model for subsequent decode step
930
- self.previous_hidden_states = HiddenStates(
931
- hidden_states, terminal_metadata,
932
- second_last_token_hidden_states)
933
- return accepted_token_ids, logprobs
934
-
935
- def _create_output_sampler_list(
936
- self,
937
- seq_group_metadata_list: List[SequenceGroupMetadata],
938
- accepted_token_ids: torch.Tensor, # shape: [batch_size, k+1]
939
- target_logprobs: torch.Tensor, # shape: [batch_size, k+1, vocab_size]
940
- prompt_logprobs: Optional[
941
- torch.Tensor], # shape: [nprompt_tokens, vocab_size]
942
- k: int,
943
- stage_times: Tuple[float, float, float],
944
- ) -> List[SamplerOutput]:
945
- """Given the accepted token ids, create a list of SamplerOutput.
946
-
947
- The output is padded with -1 tokens such that each sequence has
948
- the same number of outputs.
949
- """
950
- batch_size, num_steps = accepted_token_ids.shape
951
- accepted_token_ids_by_step = accepted_token_ids.transpose(0, 1)
952
- if self._disable_logprobs:
953
- # We are skipping the logprobs. Hence don't serialize the
954
- # logprobs related tensors from the GPU. Instead create
955
- # empty/dummy lists.
956
- (accepted_token_id_ranks_by_step,
957
- accepted_token_id_logprobs_by_step,
958
- topk_logprobs_by_step, topk_indices_by_step) =\
959
- self._create_dummy_logprob_lists(
960
- batch_size, num_steps,
961
- self.scorer_worker.model_config.max_logprobs)
962
- else:
963
- # Organize input tensors by step instead of by sequence.
964
- target_logprobs_by_step = target_logprobs.transpose(0, 1)
965
- # Serialize all tensors into Python lists.
966
- (accepted_token_id_ranks_by_step,
967
- accepted_token_id_logprobs_by_step,
968
- topk_logprobs_by_step, topk_indices_by_step) =\
969
- self._create_logprob_lists_from_tensors(
970
- target_logprobs_by_step, accepted_token_ids_by_step,
971
- self.scorer_worker.model_config.max_logprobs)
972
-
973
- # Get the sequence ids and num_logprobs (sampling parameter) in the
974
- # batch.
975
- seq_ids, request_ids_seq_ids_mapping = get_all_seq_ids_and_request_ids(
976
- seq_group_metadata_list)
977
-
978
- num_logprobs_per_seq = get_all_num_logprobs(seq_group_metadata_list)
979
-
980
- # Serialize tensor to CPU Python list.
981
- accepted_token_ids_by_step = accepted_token_ids_by_step.tolist()
982
-
983
- # Construct the output on a per-step, per-sequence basis.
984
- # Non-terminal prefill chunks will end up here as rows with just -1s
985
- # i.e mixed-batch [[-1, 1576], [-1, 29884], [-1, -1], [-1, -1]] while
986
- # terminal chunks will only have one generated token at time 0.
987
- sampler_output_list: List[SamplerOutput] = []
988
-
989
- # Prefills are not multi-step (return at most 1 token), in order to
990
- # avoid padding or repetition to fit decodes, we separate them.
991
- for i, sg in enumerate(seq_group_metadata_list):
992
- if not sg.is_prompt:
993
- # Requests are ordered as prefills|decodes=>no more prefills.
994
- break
995
- num_logprobs = num_logprobs_per_seq[i]
996
- seq_kwargs = dict(token_id=-1,
997
- token_id_logprob_rank=0,
998
- token_id_logprob=-float('inf'),
999
- topk_token_ids=[-1] * num_logprobs,
1000
- topk_logprobs=[-float('inf')] * num_logprobs,
1001
- seq_id=seq_ids[i])
1002
- # Terminal chunk, has token.
1003
- if sg.do_sample:
1004
- seq_kwargs.update(
1005
- dict(
1006
- token_id=accepted_token_ids[i][0].item(),
1007
- token_id_logprob_rank=accepted_token_id_ranks_by_step[
1008
- 0][i],
1009
- token_id_logprob=accepted_token_id_logprobs_by_step[0]
1010
- [i],
1011
- topk_token_ids=topk_indices_by_step[0][i]
1012
- [:num_logprobs],
1013
- # output only so step is 0
1014
- topk_logprobs=topk_logprobs_by_step[0][i]
1015
- [:num_logprobs],
1016
- ))
1017
- needs_plogs = (sg.sampling_params.prompt_logprobs
1018
- and sg.sampling_params.prompt_logprobs > 0)
1019
- plogs = None
1020
- if prompt_logprobs is not None:
1021
- # Even non-terminal prompt chunks can have logprobs here.
1022
- plogs = prompt_logprobs[i]
1023
- elif needs_plogs:
1024
- # Prompt logprobs are requested but `_disable_logprobs` is set.
1025
- seq_data = next(iter(sg.seq_data.values()))
1026
- # Get only the tokens in this chunk!
1027
- prompt_token_ids = seq_data.get_prompt_token_ids()
1028
- prompt_token_ids = prompt_token_ids[
1029
- seq_data.
1030
- _num_computed_tokens:seq_data._num_computed_tokens +
1031
- sg.token_chunk_size]
1032
-
1033
- is_first_chunk = seq_data._num_computed_tokens == 0
1034
- # There's no prob generated for the first token in a sequence.
1035
- if is_first_chunk:
1036
- prompt_token_ids = prompt_token_ids[1:]
1037
- plogs = [
1038
- create_logprobs_output(
1039
- token_id=p_token_id,
1040
- token_id_logprob_rank=-1,
1041
- token_id_logprob=0.0,
1042
- topk_token_ids=[],
1043
- topk_logprobs=[],
1044
- ) for p_token_id in prompt_token_ids
1045
- ]
1046
- seq_kwargs.update(dict(prompt_logprobs=plogs))
1047
-
1048
- sampler_output_list.append(
1049
- SamplerOutput(
1050
- outputs=[create_sequence_group_output(
1051
- **seq_kwargs)])) # type: ignore
1052
-
1053
- # Decodes, create one SamplerOutput per-step (at most K+1).
1054
- for step_index in range(num_steps):
1055
- if all(token_id == -1 for sg, token_id in zip(
1056
- seq_group_metadata_list,
1057
- accepted_token_ids_by_step[step_index])
1058
- if not sg.is_prompt):
1059
- break
1060
-
1061
- step_output_token_ids: List[CompletionSequenceGroupOutput] = []
1062
- for sequence_index in range(batch_size):
1063
- seq_meta = seq_group_metadata_list[sequence_index]
1064
- # Prompts already processed above.
1065
- if seq_meta.is_prompt:
1066
- continue
1067
-
1068
- # Each sequence may have a different num_logprobs; retrieve it.
1069
- num_logprobs = num_logprobs_per_seq[sequence_index]
1070
- step_output_token_ids.append(
1071
- create_sequence_group_output(
1072
- token_id=accepted_token_ids_by_step[step_index]
1073
- [sequence_index],
1074
- token_id_logprob_rank=accepted_token_id_ranks_by_step[
1075
- step_index][sequence_index],
1076
- token_id_logprob=accepted_token_id_logprobs_by_step[
1077
- step_index][sequence_index],
1078
- seq_id=seq_ids[sequence_index],
1079
- topk_token_ids=topk_indices_by_step[step_index]
1080
- [sequence_index][:num_logprobs],
1081
- topk_logprobs=topk_logprobs_by_step[step_index]
1082
- [sequence_index][:num_logprobs],
1083
- step_index=step_index))
1084
- sampler_output_list.append(
1085
- SamplerOutput(outputs=step_output_token_ids))
1086
-
1087
- # Populate the data structures needed to keep track of sequences with
1088
- # bonus tokens.
1089
- self._track_sequences_with_bonus_tokens(seq_ids,
1090
- request_ids_seq_ids_mapping,
1091
- accepted_token_ids_by_step)
1092
- maybe_rejsample_metrics = (
1093
- self._metrics.maybe_collect_rejsample_metrics(k))
1094
- if maybe_rejsample_metrics is not None:
1095
- sampler_output_list[
1096
- 0].spec_decode_worker_metrics = maybe_rejsample_metrics
1097
-
1098
- # Log time spent in each stage periodically.
1099
- # This is periodic because the rejection sampler emits metrics
1100
- # periodically.
1101
- self._maybe_log_stage_times(*stage_times)
1102
- # First `n_prefills` entries will contain prefills SamplerOutput when
1103
- # chunked prefill is enabled, the rest is decodes in multi-step format.
1104
- return sampler_output_list
1105
-
1106
- def _maybe_log_stage_times(self, average_time_per_proposal_tok_ms: float,
1107
- scoring_time_ms: float,
1108
- verification_time_ms: float) -> None:
1109
- """Log the speculative stage times. If stat logging is disabled, do
1110
- nothing.
1111
- """
1112
- if self._disable_log_stats:
1113
- return
1114
-
1115
- logger.info(
1116
- "SpecDecodeWorker stage times: "
1117
- "average_time_per_proposal_tok_ms=%.02f "
1118
- "scoring_time_ms=%.02f verification_time_ms=%.02f",
1119
- average_time_per_proposal_tok_ms, scoring_time_ms,
1120
- verification_time_ms)
1121
-
1122
- def _create_dummy_logprob_lists(
1123
- self,
1124
- batch_size: int,
1125
- num_steps: int,
1126
- num_top_k: int,
1127
- ) -> Tuple[List[List[int]], List[List[float]],
1128
- List[List[List[Optional[float]]]],
1129
- List[List[List[Optional[int]]]]]:
1130
- """
1131
- Creates and returns four dummy lists representing token probabilities
1132
- and their ranks.
1133
-
1134
- This method initializes and returns:
1135
- - The ranks of the accepted tokens, shaped (num_steps, batch_size)
1136
- - The log probabilities of the accepted tokens,
1137
- shaped (num_steps, batch_size)
1138
- - The log probabilities of the top k tokens,
1139
- shaped (num_steps, batch_size, num_top_k)
1140
- - The token IDs of the top k tokens,
1141
- shaped (num_steps, batch_size, num_top_k)
1142
-
1143
- Args:
1144
- batch_size (int): The size of the batch.
1145
- num_steps (int): The number of steps in the sequence.
1146
- num_top_k (int): The number of top-k token log probabilities to
1147
- return.
1148
-
1149
- Returns:
1150
- A tuple containing four dummy lists as described above.
1151
- """
1152
- accepted_token_id_ranks_by_step = [[-1] * batch_size
1153
- for _ in range(num_steps)]
1154
- accepted_token_id_logprobs_by_step = [[0.0] * batch_size
1155
- for _ in range(num_steps)]
1156
- topk_logprobs_by_step: List[List[List[Optional[float]]]] = [[
1157
- [None] * num_top_k for _ in range(batch_size)
1158
- ] for _ in range(num_steps)]
1159
- topk_indices_by_step: List[List[List[Optional[int]]]] = [[
1160
- [None] * num_top_k for _ in range(batch_size)
1161
- ] for _ in range(num_steps)]
1162
- return (accepted_token_id_ranks_by_step,
1163
- accepted_token_id_logprobs_by_step, topk_logprobs_by_step,
1164
- topk_indices_by_step)
1165
-
1166
- def _create_logprob_lists_from_tensors(
1167
- self,
1168
- target_logprobs_by_step: torch.Tensor,
1169
- accepted_token_ids_by_step: torch.Tensor,
1170
- num_top_k: int,
1171
- ) -> Tuple[List[List[int]], List[List[float]],
1172
- List[List[List[Optional[float]]]],
1173
- List[List[List[Optional[int]]]]]:
1174
- """
1175
- Creates and returns four lists representing token probabilities and
1176
- their ranks.
1177
-
1178
- This method initializes and returns four lists containing:
1179
- - The ranks of the accepted tokens, shaped (num_steps, batch_size)
1180
- - The log probabilities of the accepted tokens,
1181
- shaped (num_steps, batch_size)
1182
- - The log probabilities of the top k tokens,
1183
- shaped (num_steps, batch_size, num_top_k)
1184
- - The token IDs of the top k tokens,
1185
- shaped (num_steps, batch_size, num_top_k)
1186
-
1187
- Args:
1188
- target_logprobs_by_step (torch.Tensor): Tensor representing the
1189
- log probabilities of the target model,
1190
- shaped (num_steps, batch_size, vocab_size)
1191
- accepted_token_ids_by_step (torch.Tensor): Tensor representing
1192
- the accepted token_ids, shaped (num_steps, batch_size)
1193
- num_top_k (int): The number of top-k token log probabilities to
1194
- return.
1195
-
1196
- Returns:
1197
- A tuple containing the lists as described above.
1198
- """
1199
- # Serialize all tensors to CPU Python lists.
1200
- # Get the logprobs/rank of the accepted tokens.
1201
- (accepted_token_id_ranks_by_step_tensor,
1202
- accepted_token_id_logprobs_by_step_tensor
1203
- ) = get_sampled_token_logprobs(
1204
- logprob_tensor=target_logprobs_by_step,
1205
- sampled_token_ids=accepted_token_ids_by_step,
1206
- )
1207
- # Get the top-k logprobs (which may or may not include the
1208
- # logprob of the accepted token).
1209
- (topk_logprobs_by_step_tensor,
1210
- topk_indices_by_step_tensor) = target_logprobs_by_step.topk(
1211
- k=num_top_k,
1212
- dim=-1,
1213
- )
1214
- accepted_token_id_ranks_by_step = (
1215
- accepted_token_id_ranks_by_step_tensor.tolist())
1216
- accepted_token_id_logprobs_by_step = (
1217
- accepted_token_id_logprobs_by_step_tensor.tolist())
1218
- topk_logprobs_by_step = topk_logprobs_by_step_tensor.tolist()
1219
- topk_indices_by_step = topk_indices_by_step_tensor.tolist()
1220
- return (accepted_token_id_ranks_by_step,
1221
- accepted_token_id_logprobs_by_step, topk_logprobs_by_step,
1222
- topk_indices_by_step)
1223
-
1224
- def _track_finished_requests(self, execute_model_req: ExecuteModelRequest):
1225
- """
1226
- Removes the finished requests and their associated sequence ids from
1227
- internal book keeping data structures.
1228
- """
1229
- for finished_request in execute_model_req.finished_requests_ids:
1230
- for seq_id in self._request_id_seq_id_mapping[finished_request]:
1231
- self._seq_with_bonus_token_in_last_step.discard(seq_id)
1232
- del self._request_id_seq_id_mapping[finished_request]
1233
-
1234
- def _track_sequences_with_bonus_tokens(
1235
- self, seq_ids: List[int],
1236
- request_ids_seq_ids_mapping: Dict[str, Set[int]],
1237
- accepted_token_ids_by_step: List[List[int]]):
1238
- """
1239
- Updates the internal data structures which keep track of sequences
1240
- which have been assigned bonus tokens in their last forward pass.
1241
- """
1242
- for seq_index, seq_id in enumerate(seq_ids):
1243
- last_token_id = accepted_token_ids_by_step[-1][seq_index]
1244
- if last_token_id == -1:
1245
- self._seq_with_bonus_token_in_last_step.discard(seq_id)
1246
- else:
1247
- self._seq_with_bonus_token_in_last_step.add(seq_id)
1248
- for request_id, sequences in request_ids_seq_ids_mapping.items():
1249
- self._request_id_seq_id_mapping[request_id].update(sequences)
1250
-
1251
- @cached_property
1252
- def _vocab_size(self) -> int:
1253
- """Get the vocab size of the model and make sure it's consistent between
1254
- draft and target workers.
1255
- """
1256
- vocab_sizes = [
1257
- worker.vocab_size
1258
- for worker in [self.proposer_worker, self.scorer_worker]
1259
- ]
1260
- assert all(vocab_sizes[0] == vocab_size for vocab_size in vocab_sizes)
1261
- return vocab_sizes[0]
1262
-
1263
- @property
1264
- def rank(self):
1265
- return self.scorer_worker.rank
1266
-
1267
- @property
1268
- def device(self):
1269
- return self.scorer_worker.device
1270
-
1271
- @property
1272
- def _driver_rank(self) -> int:
1273
- return 0
1274
-
1275
- def get_cache_block_size_bytes(self):
1276
- """Return the size of a cache block in bytes.
1277
-
1278
- This function is only used to compose workers within a SpecDecodeWorker.
1279
- We leave composing a SpecDecodeWorker within a SpecDecodeWorker
1280
- undefined for now, although it could be implemented in the future.
1281
- See https://arxiv.org/abs/2308.04623.
1282
- """
1283
- raise NotImplementedError
1284
-
1285
- def start_profile(self):
1286
- if isinstance(self.scorer_worker, WorkerBase):
1287
- self.scorer_worker.start_profile()
1288
-
1289
- def stop_profile(self):
1290
- if isinstance(self.scorer_worker, WorkerBase):
1291
- self.scorer_worker.stop_profile()
1292
-
1293
-
1294
-def split_num_cache_blocks_evenly(scorer_cache_block_size_bytes: int,
1295
- proposer_cache_block_size_bytes: int,
1296
- total_num_gpu_blocks: int) -> int:
1297
- """Given total_num_gpu_blocks, the number of GPU blocks that could be
1298
- allocate to the target model, this function calculates how many blocks
1299
- should be given to the draft and target model.
1300
-
1301
- Note that usually the block size, in bytes, of each model is different,
1302
- as it's a function of number of KV/layer, number of heads, and hidden
1303
- dimension size.
1304
-
1305
- Since the target and draft models allocate the same number of blocks, we
1306
- simply calculate the number of blocks where if allocated by both models,
1307
- the total memory usage from KV cache is no larger than the number of
1308
- blocks allocatable by the target model alone.
1309
- """
1310
- new_num_gpu_blocks = int(
1311
- total_num_gpu_blocks * scorer_cache_block_size_bytes /
1312
- (proposer_cache_block_size_bytes + scorer_cache_block_size_bytes))
1313
-
1314
- return new_num_gpu_blocks
1315
-
1316
-
1317
-def prepare_prefill_hidden_states(
1318
- prefill_hidden_states: torch.Tensor) -> HiddenStates:
1319
- # For prefill step in proposer, we run the model for N-1 tokens
1320
- # because Nth token will be processed in the first decode step. For
1321
- # N-1 tokens, the input should be 0:N-1 hidden states which should
1322
- # be concatanated with 1:N token (since output of scorer has to be
1323
- # the input for proposer). Therefore, we shift the hidden states to
1324
- # align n-1th hidden state with nth token.
1325
- return HiddenStates(prefill_hidden_states.roll(
1326
- shifts=1, dims=0)) if prefill_hidden_states is not None else None
vllm/spec_decode/target_model_runner.py
deleted
-45
@@ -1,45 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from typing import List, Optional
5
-
6
-from vllm.sequence import SequenceGroupMetadata
7
-from vllm.worker.model_runner_base import (ModelRunnerBase,
8
- ModelRunnerInputBase,
9
- ModelRunnerWrapperBase)
10
-
11
-
12
-class TargetModelRunner(ModelRunnerWrapperBase):
13
- """Specialized model runner for speculative decoding target model.
14
- In speculative decoding, the log probabilities selected finally may not
15
- be the same ones as selected by the target model sampling. This means
16
- that the time spent in the log probability calculation of the target model
17
- is time wasted, since we calculate log probabilities after deciding which
18
- tokens are accepted. For this reason disabling log probabilities in the
19
- target model will make decode faster. The model runner sets the
20
- SamplingMetadata parameters according to whether log probabilities are
21
- requested or not.
22
- """
23
-
24
- def __init__(self, model_runner: ModelRunnerBase):
25
- # An internal boolean member variable to indicate if token log
26
- # probabilities are needed or not.
27
- super().__init__(model_runner)
28
- self.disable_logprobs = True
29
-
30
- def prepare_model_input(
31
- self,
32
- seq_group_metadata_list: List[SequenceGroupMetadata],
33
- virtual_engine: int = 0,
34
- finished_requests_ids: Optional[List[str]] = None,
35
- ) -> ModelRunnerInputBase:
36
- model_input: ModelRunnerInputBase =\
37
- self.model_runner.prepare_model_input(
38
- seq_group_metadata_list, virtual_engine, finished_requests_ids)
39
- # If token log probabilities is disabled then skip generating sampler
40
- # CPU output. We directly serialize the GPU sampled_token_id tensors
41
- # as needed. If log probabilities is enabled then synchronize all the
42
- # sampling related tensors which includes the logprobs tensors.
43
- model_input.sampling_metadata.skip_sampler_cpu_output = (
44
- self.disable_logprobs)
45
- return model_input
vllm/spec_decode/top1_proposer.py
deleted
-275
@@ -1,275 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from typing import List, Optional, Set, Tuple
5
-
6
-import torch
7
-
8
-from vllm.model_executor.layers.sampler import SamplerOutput
9
-from vllm.sequence import ExecuteModelRequest, SequenceGroupMetadata
10
-from vllm.spec_decode.interfaces import (SpeculativeProposals,
11
- SpeculativeProposer)
12
-from vllm.spec_decode.proposer_worker_base import ProposerWorkerBase
13
-from vllm.spec_decode.util import sampler_output_to_torch
14
-
15
-
16
-class Top1Proposer(SpeculativeProposer):
17
- """Helper class which separates out sequences which would exceed the max
18
- model length when speculated upon.
19
-
20
- This allows combinations of models such as JackFram/llama-68m draft with
21
- meta-llama/Llama2-13b-chat-hf, as llama-68m has max_position_embeddings of
22
- 2048 while Llama2-13b has max_position_embeddings of 4096.
23
-
24
- We treat the sequences which exceed the proposal draft model length as
25
- "non-spec sequences". Essentially they skip the draft model and go through
26
- normal decoding in the target model.
27
-
28
- Currently, only proposal_lens of 0 and k are supported, where k is a global
29
- batch proposal length. In the future vLLM should support per-sequence
30
- proposal lengths.
31
- """
32
-
33
- def __init__(
34
- self,
35
- worker: ProposerWorkerBase,
36
- device: str,
37
- vocab_size: int,
38
- max_proposal_len: Optional[int] = None,
39
- ):
40
- self._worker = worker
41
- self._device = device
42
- self.max_proposal_len = max_proposal_len
43
- self._vocab_size = vocab_size
44
-
45
- def get_spec_proposals(
46
- self,
47
- execute_model_req: ExecuteModelRequest,
48
- seq_ids_with_bonus_token_in_last_step: Set[int],
49
- ) -> SpeculativeProposals:
50
- """Get speculative proposals given the input batch.
51
-
52
- Sequences which would exceed the max model length are skipped during
53
- speculation.
54
- """
55
- proposal_len = execute_model_req.num_lookahead_slots
56
- seq_group_metadata_list = execute_model_req.seq_group_metadata_list
57
-
58
- # Split speculative- and non-speculative- sequences.
59
- (
60
- proposal_lens,
61
- nonzero_proposal_len_seqs,
62
- nonzero_proposal_len_indices,
63
- ) = self._split_by_proposal_len(seq_group_metadata_list, proposal_len)
64
-
65
- if nonzero_proposal_len_seqs:
66
- # Speculate tokens using the draft worker for the speculative
67
- # sequences.
68
- # If sampler_transposed is true, then maybe_sampler_output's
69
- # token_ids is like [batch] format in proposal_len size list,
70
- # while if it is false, the format would be [proposal_len]
71
- # in batch size list
72
- hidden_states = execute_model_req.previous_hidden_states
73
- if hidden_states is not None:
74
- hidden_states.prune(nonzero_proposal_len_seqs)
75
- nonzero_execute_model_req = ExecuteModelRequest(
76
- seq_group_metadata_list=nonzero_proposal_len_seqs,
77
- num_lookahead_slots=proposal_len,
78
- previous_hidden_states=hidden_states,
79
- )
80
- maybe_sampler_output, transposed = self._worker.sampler_output(
81
- execute_model_req=nonzero_execute_model_req,
82
- sample_len=proposal_len,
83
- seq_ids_with_bonus_token_in_last_step=\
84
- seq_ids_with_bonus_token_in_last_step,
85
- )
86
- (
87
- proposal_lens,
88
- maybe_sampler_output,
89
- nonzero_proposal_len_indices,
90
- ) = self._remove_no_proposal_seqs(proposal_lens,
91
- maybe_sampler_output,
92
- nonzero_proposal_len_indices,
93
- transposed)
94
- else:
95
- # If no sequences can be speculated, set sampler output to None.
96
- maybe_sampler_output = None
97
- transposed = False
98
-
99
- # Combine speculative- and non-speculative sequences into the same
100
- # representation.
101
- proposal_tokens, proposal_probs, proposal_lens = self._merge_outputs(
102
- batch_size=len(seq_group_metadata_list),
103
- proposal_len=proposal_len,
104
- maybe_sampler_output=maybe_sampler_output,
105
- proposal_lens=proposal_lens,
106
- nonzero_proposal_len_indices=nonzero_proposal_len_indices,
107
- sampler_transposed=transposed,
108
- )
109
-
110
- proposals = SpeculativeProposals(proposal_token_ids=proposal_tokens,
111
- proposal_probs=proposal_probs,
112
- proposal_lens=proposal_lens,
113
- no_proposals=maybe_sampler_output
114
- is None)
115
- return proposals
116
-
117
- def _split_by_proposal_len(
118
- self,
119
- seq_group_metadata_list: List[SequenceGroupMetadata],
120
- proposal_len: int,
121
- ) -> Tuple[List[int], List[SequenceGroupMetadata], List[int]]:
122
- """Split sequences by two groups:
123
- 1. Sequences with non-zero proposal length.
124
- 2. Sequences with zero proposal length (due to disabled speculation
125
- or exceed the maximum model length).
126
- """
127
-
128
- proposal_lens: List[int] = []
129
- nonzero_proposal_len_seqs: List[SequenceGroupMetadata] = []
130
- nonzero_proposal_len_indices: List[int] = []
131
- for i, seq_group_metadata in enumerate(seq_group_metadata_list):
132
- # The speculative decoding for this request has either been disabled
133
- # (e.g. due to high traffic) or this is a prompt request.
134
- if (seq_group_metadata.is_prompt
135
- or seq_group_metadata.num_speculative_tokens == 0):
136
- proposal_lens.append(0)
137
- continue
138
-
139
- seq_data = next(iter(seq_group_metadata.seq_data.values()))
140
- seq_len = seq_data.get_len()
141
-
142
- # Currently only proposal lens of 0 or the global batch proposal len
143
- # are supported.
144
- # If max_proposal_len is defined, then we shall not exceed this
145
- # quota for nonzero_proposal
146
- new_k = 0
147
- if (self.max_proposal_len is None
148
- or seq_len + proposal_len < self.max_proposal_len):
149
- new_k = proposal_len
150
- nonzero_proposal_len_seqs.append(seq_group_metadata)
151
- nonzero_proposal_len_indices.append(i)
152
- proposal_lens.append(new_k)
153
- seq_group_metadata.num_speculative_tokens = new_k
154
-
155
- return (
156
- proposal_lens,
157
- nonzero_proposal_len_seqs,
158
- nonzero_proposal_len_indices,
159
- )
160
-
161
- @staticmethod
162
- def _remove_no_proposal_seqs(proposal_lens, maybe_sampler_output,
163
- nonzero_proposal_len_indices, transposed):
164
- """Remove sequences from nonzero_proposal_len_indices and reset
165
- their proposal_len to 0 the draft worker does not provide a proposal
166
- (maybe_sampler_output=None). This can avoid scoring overheads.
167
- """
168
-
169
- # If maybe_sampler_output is None, then the draft worker did not
170
- # provide a proposal for any sequence and thus no action needed.
171
- # Also we do not support transposed maybe_sampler_output for now
172
- # because it seems not straightforward for draft workers outputting
173
- # transposed sampler outputs to handle the case of no proposal.
174
- if maybe_sampler_output is None or transposed:
175
- return (proposal_lens, maybe_sampler_output,
176
- nonzero_proposal_len_indices)
177
-
178
- new_proposal_lens: List[int] = []
179
- new_nonzero_proposal_len_indices: List[int] = []
180
- new_maybe_sampler_output: List[SamplerOutput] = []
181
- nonzero_proposal_len_idx_ptr = 0
182
- seq_idx = 0
183
- while seq_idx < len(
184
- proposal_lens) and nonzero_proposal_len_idx_ptr < len(
185
- nonzero_proposal_len_indices):
186
- if seq_idx < nonzero_proposal_len_indices[
187
- nonzero_proposal_len_idx_ptr]:
188
- # Sequence is not in the original nonzero_proposal_len_indices,
189
- # meaning that it has a proposal length of 0 before sending to
190
- # the draft worker.
191
- assert proposal_lens[seq_idx] == 0
192
- new_proposal_lens.append(0)
193
- else:
194
- # Sequence is in the original nonzero_proposal_len_indices
195
- if maybe_sampler_output[nonzero_proposal_len_idx_ptr] is None:
196
- # but does not have a proposal from the draft worker.
197
- new_proposal_lens.append(0)
198
- else:
199
- # and has a proposal from the draft worker. Add it to the
200
- # new nonzero proposal list and keep the sampler output.
201
- new_proposal_lens.append(proposal_lens[seq_idx])
202
- new_nonzero_proposal_len_indices.append(seq_idx)
203
- new_maybe_sampler_output.append(
204
- maybe_sampler_output[nonzero_proposal_len_idx_ptr])
205
- nonzero_proposal_len_idx_ptr += 1
206
- seq_idx += 1
207
-
208
- # The remaining sequences should have proposal length of 0.
209
- new_proposal_lens.extend(proposal_lens[seq_idx:])
210
-
211
- # We assume sampler_output will not be a list of all Nones.
212
- # In this case this function should not be called.
213
- assert new_maybe_sampler_output
214
- return (new_proposal_lens, new_maybe_sampler_output,
215
- new_nonzero_proposal_len_indices)
216
-
217
- def _merge_outputs(
218
- self,
219
- batch_size: int,
220
- proposal_len: int,
221
- maybe_sampler_output: Optional[List[SamplerOutput]],
222
- proposal_lens: List[int],
223
- nonzero_proposal_len_indices: List[int],
224
- sampler_transposed: bool,
225
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
226
- """After speculations are produced, merge the speculation results with
227
- the skipped sequences.
228
- """
229
- if maybe_sampler_output is None:
230
- # If no speculative tokens, the sampler output will be None.
231
- # In this case we return empty proposals.
232
- proposal_tokens = torch.tensor(-1,
233
- dtype=torch.long,
234
- device=self._device).expand(
235
- batch_size, proposal_len)
236
- proposal_probs = torch.tensor(0,
237
- dtype=torch.float32,
238
- device=self._device).expand(
239
- batch_size, proposal_len,
240
- self._vocab_size)
241
- proposal_lens_tensor = torch.tensor(0,
242
- dtype=torch.long,
243
- device=self._device).expand(
244
- len(proposal_lens))
245
- return proposal_tokens, proposal_probs, proposal_lens_tensor
246
-
247
- sampler_output = maybe_sampler_output
248
- proposal_tokens, proposal_probs, *_ = sampler_output_to_torch(
249
- sampler_output, sampler_transposed)
250
-
251
- # Now, reformat the output GPU tensors such that each sequence has
252
- # a proposal. the proposal can be empty, e.g. [-1, -1, -1]
253
-
254
- entire_proposal_tokens = proposal_tokens.new_full(
255
- size=(batch_size, *proposal_tokens.shape[1:]),
256
- fill_value=-1,
257
- )
258
- entire_proposal_tokens[nonzero_proposal_len_indices] = proposal_tokens
259
- entire_proposal_probs = proposal_probs.new_zeros(
260
- batch_size,
261
- *proposal_probs.shape[1:],
262
- )
263
- entire_proposal_probs[nonzero_proposal_len_indices] = proposal_probs
264
-
265
- proposal_tokens, proposal_probs = (
266
- entire_proposal_tokens,
267
- entire_proposal_probs,
268
- )
269
-
270
- proposal_lens_tensor = torch.zeros(batch_size,
271
- dtype=torch.long,
272
- device=self._device)
273
- proposal_lens_tensor[nonzero_proposal_len_indices] = proposal_len
274
-
275
- return proposal_tokens, proposal_probs, proposal_lens_tensor
vllm/spec_decode/util.py
deleted
-277
@@ -1,277 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import time
5
-from contextlib import contextmanager
6
-from typing import Dict, List, Optional, Sequence, Tuple
7
-
8
-import torch
9
-
10
-from vllm.model_executor.layers.sampler import SamplerOutput
11
-from vllm.platforms import current_platform
12
-from vllm.sequence import (CompletionSequenceGroupOutput, Logprob,
13
- PromptLogprobs, SequenceGroupMetadata,
14
- SequenceOutput)
15
-
16
-SeqId = int
17
-
18
-
19
-def get_all_num_logprobs(
20
- seq_group_metadata_list: List[SequenceGroupMetadata]) -> List[int]:
21
- """Given a list of SequenceGroupMetadata, create a list of all num_logprobs.
22
-
23
- If the sampling params do not call for any logprobs, return 0 for that
24
- sequence.
25
- """
26
-
27
- all_num_logprobs: List[int] = []
28
- for seq_group_metadata in seq_group_metadata_list:
29
- num_logprobs = seq_group_metadata.sampling_params.logprobs
30
- if num_logprobs is None:
31
- num_logprobs = 0
32
- all_num_logprobs.append(num_logprobs)
33
-
34
- return all_num_logprobs
35
-
36
-
37
-def get_sampled_token_logprobs(
38
- # shape [num_steps, batch_size, vocab_size]
39
- logprob_tensor: torch.Tensor,
40
- sampled_token_ids: torch.Tensor, # shape [num_steps, batch_size]
41
-) -> Tuple[torch.Tensor, torch.Tensor]:
42
- """Get the logprobs for the sampled tokens. Returns the ranks and logprobs.
43
- """
44
- num_steps, batch_size, vocab_size = logprob_tensor.shape
45
-
46
- selected_logprobs = logprob_tensor[
47
- torch.arange(num_steps).unsqueeze(1),
48
- torch.arange(batch_size),
49
- sampled_token_ids,
50
- ]
51
- expanded_selected_logprobs = selected_logprobs.unsqueeze(-1).expand(
52
- -1, -1, vocab_size)
53
- sampled_token_ids_ranks = (logprob_tensor
54
- > expanded_selected_logprobs).sum(-1).add_(1)
55
-
56
- return sampled_token_ids_ranks, selected_logprobs
57
-
58
-
59
-def create_logprobs_output(
60
- token_id: int,
61
- token_id_logprob_rank: int,
62
- token_id_logprob: float,
63
- topk_token_ids: List[Optional[int]],
64
- topk_logprobs: List[Optional[float]],
65
-) -> Dict[int, Logprob]:
66
- """Create a Logprob Dict for a token given the sampling results.
67
-
68
- Args:
69
- token_id (int): The sampled token for the sequence.
70
- token_id_logprob_rank (int): The logprob rank of the sampled token.
71
- token_id_logprob (float): The logprob value of the sampled token.
72
- topk_token_ids (List[Optional[int]]): The list of top-k token ids.
73
- topk_logprobs (List[Optional[float]]): The list of top-k logprobs.
74
- """
75
- # vLLM logprobs always include the sampled token. In addition, the user may
76
- # request topk-logprobs (where top-k varies per user up to max_logprobs).
77
- logprobs: Dict[int, Logprob] = {
78
- token_id: Logprob(
79
- logprob=token_id_logprob,
80
- rank=token_id_logprob_rank,
81
- ),
82
- }
83
- logprobs.update({
84
- topk_token_id: Logprob(
85
- logprob=topk_logprob if topk_logprob is not None else 0.0,
86
- rank=topk_index + 1,
87
- )
88
- for topk_index, (topk_token_id, topk_logprob) \
89
- in enumerate(zip(topk_token_ids, topk_logprobs)) \
90
- if topk_token_id is not None
91
- })
92
-
93
- return logprobs
94
-
95
-
96
-def create_sequence_group_output(
97
- token_id: int,
98
- token_id_logprob_rank: int,
99
- token_id_logprob: float,
100
- seq_id: SeqId,
101
- topk_token_ids: List[Optional[int]],
102
- topk_logprobs: List[Optional[float]],
103
- prompt_logprobs: Optional[PromptLogprobs] = None,
104
- step_index: Optional[int] = 0) -> CompletionSequenceGroupOutput:
105
- """Create a SequenceGroupOutput given the sampling results.
106
-
107
- Args:
108
- token_id (int): The sampled token for the sequence.
109
- token_id_logprob_rank (int): The logprob rank of the sampled token.
110
- token_id_logprob (float): The logprob value of the sampled token.
111
- seq_id (int): The sequence id.
112
- topk_token_ids (List[Optional[int]]): The list of top-k token ids.
113
- topk_logprobs (List[Optional[float]]): The list of top-k logprobs.
114
- step_index: (Optional[int]): The index of the speculative token.
115
- """
116
-
117
- logprobs = create_logprobs_output(
118
- token_id,
119
- token_id_logprob_rank,
120
- token_id_logprob,
121
- topk_token_ids,
122
- topk_logprobs,
123
- )
124
-
125
- return CompletionSequenceGroupOutput(samples=[
126
- SequenceOutput(parent_seq_id=seq_id,
127
- output_token=token_id,
128
- logprobs=logprobs)
129
- ],
130
- prompt_logprobs=prompt_logprobs,
131
- step_index=step_index)
132
-
133
-
134
-def split_batch_by_proposal_len(
135
- seq_group_metadata_list: List[SequenceGroupMetadata],
136
- proposal_lens: List[int],
137
-) -> Tuple[Tuple[List[SequenceGroupMetadata], List[int]], Tuple[
138
- List[SequenceGroupMetadata], List[int]]]:
139
- """Utility function that splits a batch based on whether the proposal len is
140
- zero or not. We should remove this once vLLM supports per-sequence proposal
141
- lens in a batch.
142
- """
143
-
144
- nonzero_lists: Tuple[List[SequenceGroupMetadata], List[int]] = ([], [])
145
- zero_lists: Tuple[List[SequenceGroupMetadata], List[int]] = ([], [])
146
- for i, (seq_group, proposal_len) in enumerate(
147
- zip(seq_group_metadata_list, proposal_lens)):
148
- seq_groups, indices = nonzero_lists if proposal_len else zero_lists
149
- seq_groups.append(seq_group)
150
- indices.append(i)
151
- return nonzero_lists, zero_lists
152
-
153
-
154
-def sampler_output_to_torch(
155
- sampler_output_list: Sequence[SamplerOutput], sampler_transposed: bool
156
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
157
- """Utility function which converts a list of SamplerOutput to tensors.
158
-
159
- sampler_transposed here is used as the indicator for whether
160
- we need do additional tensor transpose logic here.
161
-
162
- Returns:
163
- sampled_token_ids: torch.Tensor
164
- shape: [batch_size, len(sampler_output_list)]
165
-
166
- sampled_token_probs: torch.Tensor
167
- shape: [batch_size, len(sampler_output_list), vocab_size]
168
- """
169
-
170
- # shape: [batch_size, num_sampler_output, vocab_size]
171
- sampled_token_probs = torch.stack(
172
- [
173
- sampler_output.sampled_token_probs
174
- for sampler_output in sampler_output_list
175
- ],
176
- dim=0,
177
- )
178
-
179
- # shape: [batch_size, num_sampler_output, vocab_size]
180
- sampled_token_logprobs = torch.stack(
181
- [sampler_output.logprobs for sampler_output in sampler_output_list],
182
- dim=0,
183
- )
184
-
185
- # shape: [batch_size, num_sampler_output]
186
- sampled_token_ids = torch.stack(
187
- [
188
- sampler_output.sampled_token_ids.flatten()
189
- for sampler_output in sampler_output_list
190
- ],
191
- dim=0,
192
- )
193
-
194
- if sampler_transposed:
195
- sampled_token_probs = sampled_token_probs.transpose(0, 1)
196
- sampled_token_logprobs = sampled_token_logprobs.transpose(0, 1)
197
- sampled_token_ids = sampled_token_ids.transpose(0, 1)
198
-
199
- if sampler_output_list[0].hidden_states is not None:
200
- # shape: [batch_size, num_sampler_output, hidden_dim]
201
- sampled_hidden_states = torch.stack(
202
- [
203
- sampler_output.hidden_states
204
- for sampler_output in sampler_output_list
205
- ],
206
- dim=0,
207
- )
208
-
209
- if sampler_transposed:
210
- sampled_hidden_states = sampled_hidden_states.transpose(0, 1)
211
- else:
212
- sampled_hidden_states = None
213
-
214
- return (sampled_token_ids, sampled_token_probs, sampled_token_logprobs,
215
- sampled_hidden_states)
216
-
217
-
218
-def maybe_mock_device_tensors(sampler_output: SamplerOutput, batch_size: int,
219
- vocab_size: int, device: str) -> None:
220
- """Helper method which mocks out the GPU tensors in SamplerOutput with dummy
221
- values. This will be removed in PR 7/9.
222
- https://docs.google.com/document/d/1rE4pr3IdspRw97XbImY4fS9IWYuJJ3HGtL7AdIKGrw8/edit#heading=h.qijw1sdidrer
223
- """
224
- values = [
225
- sampler_output.sampled_token_probs, sampler_output.sampled_token_ids
226
- ]
227
- assert all(v is None for v in values) or not any(v is None for v in values)
228
- if not any(v is None for v in values):
229
- # Do nothing if the tensors are already created (usually in unit tests).
230
- return
231
-
232
- # Softmax to ensure valid probs.
233
- sampler_output.sampled_token_probs = torch.nn.functional.softmax(
234
- torch.rand(batch_size, vocab_size, dtype=torch.float32, device=device),
235
- dim=-1)
236
-
237
- sampler_output.sampled_token_ids = torch.randint(low=10,
238
- high=100,
239
- size=(batch_size, ),
240
- dtype=torch.long,
241
- device=device)
242
-
243
-
244
-@contextmanager
245
-def nvtx_range(msg, *args, **kwargs):
246
- """
247
- Context manager / decorator that pushes an NVTX range at the beginning
248
- of its scope, and pops it at the end. If extra arguments are given,
249
- they are passed as arguments to msg.format().
250
-
251
- If running with cuda graphs, you must enable nsys cuda graph profiling.
252
-
253
- Arguments:
254
- msg (string): message to associate with the range
255
- """
256
- if current_platform.is_cuda_alike():
257
- torch.cuda.nvtx.range_push(msg.format(*args, **kwargs))
258
- try:
259
- yield
260
- finally:
261
- torch.cuda.nvtx.range_pop()
262
- else:
263
- yield
264
-
265
-
266
-class Timer:
267
- """Basic timer context manager for measuring CPU time.
268
- """
269
-
270
- def __enter__(self):
271
- self.start_time = time.time()
272
- return self
273
-
274
- def __exit__(self, exc_type, exc_value, traceback):
275
- self.end_time = time.time()
276
- self.elapsed_time_s = self.end_time - self.start_time
277
- self.elapsed_time_ms = self.elapsed_time_s * 1000
vllm/transformers_utils/configs/eagle.py
+18
-22
@@ -6,7 +6,6 @@ from typing import Optional, Union
6
7
from transformers import AutoConfig, PretrainedConfig
8
9
-import vllm.envs as envs
9
from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekV2Config
10
11
@@ -44,28 +43,25 @@ class EAGLEConfig(PretrainedConfig):
43
self.truncated_vocab_size = self.model.vocab_size if \
44
truncated_vocab_size is None else truncated_vocab_size
45
47
- if not envs.VLLM_USE_V1:
48
- kwargs["architectures"] = ["EAGLEModel"]
46
+ # Eagle model name should follow naming convention of
47
+ # LlamaForCausalLM -> EagleLlamaForCausalLM
48
+ if method == "eagle":
49
+ assert self.model is not None, \
50
+ "model should not be None when method is eagle"
51
+ kwargs["architectures"] = [
52
+ f"Eagle{arch}" if not arch.startswith("Eagle") \
53
+ else arch for arch in self.model.architectures
54
+ ]
55
+ elif method == "eagle3":
56
+ assert self.model is not None, \
57
+ "model should not be None when method is eagle3"
58
+ kwargs["architectures"] = [
59
+ f"Eagle3{arch}" if not arch.startswith("Eagle3") \
60
+ else arch for arch in self.model.architectures
61
+ ]
62
else:
50
- # Eagle model name should follow naming convention of
51
- # LlamaForCausalLM -> EagleLlamaForCausalLM
52
- if method == "eagle":
53
- assert self.model is not None, \
54
- "model should not be None when method is eagle"
55
- kwargs["architectures"] = [
56
- f"Eagle{arch}" if not arch.startswith("Eagle") \
57
- else arch for arch in self.model.architectures
58
- ]
59
- elif method == "eagle3":
60
- assert self.model is not None, \
61
- "model should not be None when method is eagle3"
62
- kwargs["architectures"] = [
63
- f"Eagle3{arch}" if not arch.startswith("Eagle3") \
64
- else arch for arch in self.model.architectures
65
- ]
66
- else:
67
- raise ValueError(f"Invalid method {method}. \
68
- Supported methods are eagle and eagle3.")
63
+ raise ValueError(f"Invalid method {method}. \
64
+ Supported methods are eagle and eagle3.")
65
66
super().__init__(**kwargs)
67
vllm/worker/worker_base.py
-2
@@ -397,8 +397,6 @@ class LocalOrDistributedWorkerBase(WorkerBase):
397
398
model_input, worker_input, kwargs = inputs
399
num_steps = worker_input.num_steps
400
- if execute_model_req is not None and execute_model_req.spec_step_idx:
401
- kwargs["spec_step_idx"] = execute_model_req.spec_step_idx
400
401
self.execute_worker(worker_input)
402