| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import time |
| 5 | |
| 6 | from vllm import LLM, SamplingParams |
| 7 | |
| 8 | # Sample prompts. |
| 9 | prompts = [ |
| 10 | "Hello, my name is", |
| 11 | "The president of the United States is", |
| 12 | "The capital of France is", |
| 13 | "The future of AI is", |
| 14 | ] |
| 15 | # Create a sampling params object. |
| 16 | sampling_params = SamplingParams(temperature=0.8, top_p=0.95) |
| 17 | |
| 18 | |
| 19 | def main(): |
| 20 | # Create an LLM. |
| 21 | llm = LLM( |
| 22 | model="facebook/opt-125m", |
| 23 | tensor_parallel_size=1, |
| 24 | profiler_config={ |
| 25 | "profiler": "torch", |
| 26 | "torch_profiler_dir": "./vllm_profile", |
| 27 | }, |
| 28 | ) |
| 29 | |
| 30 | llm.start_profile() |
| 31 | |
| 32 | # Generate texts from the prompts. The output is a list of RequestOutput |
| 33 | # objects that contain the prompt, generated text, and other information. |
| 34 | outputs = llm.generate(prompts, sampling_params) |
| 35 | |
| 36 | llm.stop_profile() |
| 37 | |
| 38 | # Print the outputs. |
| 39 | print("-" * 50) |
| 40 | for output in outputs: |
| 41 | prompt = output.prompt |
| 42 | generated_text = output.outputs[0].text |
| 43 | print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}") |
| 44 | print("-" * 50) |
| 45 | |
| 46 | # Add a buffer to wait for profiler in the background process |
| 47 | # (in case MP is on) to finish writing profiling output. |
| 48 | time.sleep(10) |
| 49 | |
| 50 | |
| 51 | if __name__ == "__main__": |
| 52 | main() |