Gahow Wang
445e491123
third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:
vllm/v1/core/sched/scheduler.py:
Replace fatal assert with graceful skip when KV transfer callback
arrives for an already-aborted request during PD disaggregated serving.
Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
75 lines
2.3 KiB
Python
75 lines
2.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
This file demonstrates using the `LLMEngine`
|
|
for processing prompts with various sampling parameters.
|
|
"""
|
|
|
|
import argparse
|
|
|
|
from vllm import EngineArgs, LLMEngine, RequestOutput, SamplingParams
|
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
|
|
|
|
|
def create_test_prompts() -> list[tuple[str, SamplingParams]]:
|
|
"""Create a list of test prompts with their sampling parameters."""
|
|
return [
|
|
(
|
|
"A robot may not injure a human being",
|
|
SamplingParams(temperature=0.0, logprobs=1, prompt_logprobs=1),
|
|
),
|
|
(
|
|
"To be or not to be,",
|
|
SamplingParams(temperature=0.8, top_k=5, presence_penalty=0.2),
|
|
),
|
|
(
|
|
"What is the meaning of life?",
|
|
SamplingParams(n=2, temperature=0.8, top_p=0.95, frequency_penalty=0.1),
|
|
),
|
|
]
|
|
|
|
|
|
def process_requests(engine: LLMEngine, test_prompts: list[tuple[str, SamplingParams]]):
|
|
"""Continuously process a list of prompts and handle the outputs."""
|
|
request_id = 0
|
|
|
|
print("-" * 50)
|
|
while test_prompts or engine.has_unfinished_requests():
|
|
if test_prompts:
|
|
prompt, sampling_params = test_prompts.pop(0)
|
|
engine.add_request(str(request_id), prompt, sampling_params)
|
|
request_id += 1
|
|
|
|
request_outputs: list[RequestOutput] = engine.step()
|
|
|
|
for request_output in request_outputs:
|
|
if request_output.finished:
|
|
print(request_output)
|
|
print("-" * 50)
|
|
|
|
|
|
def initialize_engine(args: argparse.Namespace) -> LLMEngine:
|
|
"""Initialize the LLMEngine from the command line arguments."""
|
|
engine_args = EngineArgs.from_cli_args(args)
|
|
return LLMEngine.from_engine_args(engine_args)
|
|
|
|
|
|
def parse_args():
|
|
parser = FlexibleArgumentParser(
|
|
description="Demo on using the LLMEngine class directly"
|
|
)
|
|
parser = EngineArgs.add_cli_args(parser)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main(args: argparse.Namespace):
|
|
"""Main function that sets up and runs the prompt processing."""
|
|
engine = initialize_engine(args)
|
|
test_prompts = create_test_prompts()
|
|
process_requests(engine, test_prompts)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
main(args)
|