first working attempt at logprobs

Author: afeldman-nm

examples/offline_inference_multi_step.py

@@ -25,6 +25,7 @@
     gpu_memory_utilization=0.9,
     num_scheduler_steps=8,
     use_v2_block_manager=True,
+    enforce_eager=True,
 )
 # Generate texts from the prompts. The output is a list of RequestOutput objects
 # that contain the prompt, generated text, and other information.

tests/spec_decode/test_multi_step_worker.py

@@ -5,8 +5,9 @@
 import pytest
 import torch
 
+from vllm.model_executor.layers.sampler import SamplerOutput
 from vllm.model_executor.utils import set_random_seed
-from vllm.sequence import ExecuteModelRequest, Logprob, SamplerOutput
+from vllm.sequence import ExecuteModelRequest, Logprob
 from vllm.spec_decode.draft_model_runner import TP1DraftModelRunner
 from vllm.spec_decode.multi_step_worker import MultiStepWorker
 from vllm.spec_decode.top1_proposer import Top1Proposer

tests/spec_decode/test_spec_decode_worker.py

@@ -7,8 +7,9 @@
 import pytest
 import torch
 
+from vllm.model_executor.layers.sampler import SamplerOutput
 from vllm.model_executor.utils import set_random_seed
-from vllm.sequence import ExecuteModelRequest, SamplerOutput, SequenceOutput
+from vllm.sequence import ExecuteModelRequest, SequenceOutput
 from vllm.spec_decode.interfaces import SpeculativeProposals
 from vllm.spec_decode.metrics import (AsyncMetricsCollector,
                                       SpecDecodeWorkerMetrics)

tests/spec_decode/utils.py

@@ -8,12 +8,12 @@
 import torch
 
 from vllm.engine.arg_utils import EngineArgs
+from vllm.model_executor.layers.sampler import SamplerOutput
 from vllm.model_executor.utils import set_random_seed
 from vllm.sampling_params import SamplingParams
 from vllm.sequence import (VLLM_TOKEN_ID_ARRAY_TYPE,
                            CompletionSequenceGroupOutput, Logprob,
-                           SamplerOutput, SequenceData, SequenceGroupMetadata,
-                           SequenceOutput)
+                           SequenceData, SequenceGroupMetadata, SequenceOutput)
 from vllm.utils import get_distributed_init_method, get_ip, get_open_port
 from vllm.worker.cache_engine import CacheEngine
 from vllm.worker.model_runner import ModelRunner

tests/test_sequence.py

@@ -2,9 +2,10 @@
 
 import pytest
 
+from vllm.model_executor.layers.sampler import SamplerOutput
 from vllm.sequence import (VLLM_TOKEN_ID_ARRAY_TYPE,
-                           CompletionSequenceGroupOutput, SamplerOutput,
-                           SequenceData, SequenceOutput)
+                           CompletionSequenceGroupOutput, SequenceData,
+                           SequenceOutput)
 
 from .core.utils import create_dummy_prompt
 

vllm/engine/async_llm_engine.py

@@ -25,12 +25,12 @@
 from vllm.inputs.parse import is_explicit_encoder_decoder_prompt
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
+from vllm.model_executor.layers.sampler import SamplerOutput
 from vllm.outputs import EmbeddingRequestOutput, RequestOutput
 from vllm.pooling_params import PoolingParams
 from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sampling_params import SamplingParams
-from vllm.sequence import (ExecuteModelRequest, SamplerOutput,
-                           SequenceGroupMetadata)
+from vllm.sequence import ExecuteModelRequest, SequenceGroupMetadata
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils import print_warning_once
 

vllm/engine/llm_engine.py

@@ -29,16 +29,16 @@
 from vllm.inputs.parse import is_explicit_encoder_decoder_prompt
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
+from vllm.model_executor.layers.sampler import SamplerOutput
 from vllm.multimodal import MultiModalDataDict
 from vllm.outputs import (EmbeddingRequestOutput, RequestOutput,
                           RequestOutputFactory)
 from vllm.pooling_params import PoolingParams
 from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sampling_params import SamplingParams
 from vllm.sequence import (EmbeddingSequenceGroupOutput, ExecuteModelRequest,
-                           PoolerOutput, SamplerOutput, Sequence,
-                           SequenceGroup, SequenceGroupMetadata,
-                           SequenceStatus)
+                           PoolerOutput, Sequence, SequenceGroup,
+                           SequenceGroupMetadata, SequenceStatus)
 from vllm.tracing import (SpanAttributes, SpanKind, extract_trace_context,
                           init_tracer)
 from vllm.transformers_utils.config import try_get_generation_config

vllm/engine/output_processor/util.py

@@ -2,7 +2,8 @@
 from typing import Sequence as GenericSequence
 from typing import Union
 
-from vllm.sequence import PoolerOutput, SamplerOutput, SequenceGroupOutput
+from vllm.model_executor.layers.sampler import SamplerOutput
+from vllm.sequence import PoolerOutput, SequenceGroupOutput
 
 
 def create_output_by_sequence_group(

vllm/engine/protocol.py

@@ -7,11 +7,11 @@
 from vllm.core.scheduler import SchedulerOutputs
 from vllm.inputs.data import PromptInputs
 from vllm.lora.request import LoRARequest
+from vllm.model_executor.layers.sampler import SamplerOutput
 from vllm.outputs import EmbeddingRequestOutput, RequestOutput
 from vllm.pooling_params import PoolingParams
 from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sampling_params import SamplingParams
-from vllm.sequence import SamplerOutput
 
 
 @runtime_checkable

vllm/executor/cpu_executor.py

@@ -11,8 +11,9 @@
                                                   ResultHandler, WorkerMonitor)
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
+from vllm.model_executor.layers.sampler import SamplerOutput
 from vllm.prompt_adapter.request import PromptAdapterRequest
-from vllm.sequence import ExecuteModelRequest, SamplerOutput
+from vllm.sequence import ExecuteModelRequest
 from vllm.utils import (GiB_bytes, get_distributed_init_method, get_open_port,
                         get_vllm_instance_id, make_async)
 from vllm.worker.worker_base import WorkerWrapperBase