sample code for vllm local inference

tests/vllm-inference-test.py

@@ -0,0 +1,86 @@
+import asyncio
+import time
+
+from vllm import LLM, SamplingParams
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.engine.async_llm_engine import AsyncLLMEngine
+from vllm.utils import random_uuid
+
+sampling_params = SamplingParams(
+    temperature=0.8,
+    top_p=0.95,
+    max_tokens=4096
+)
+
+prompt = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are a helpful LLM in a WebRTC call. Your goal is to demonstrate your capabilities in a succinct way. Your output will be converted to audio so don't include special characters in your answers. Respond to what the user said in a creative and helpful way.<|eot_id|><|start_header_id|>system<|end_header_id|>\n\nPlease introduce yourself to the user.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+
+
+async def main():
+    print("🥶 cold starting inference")
+    start = time.monotonic_ns()
+
+    engine_args = AsyncEngineArgs(
+        model="meta-llama/Meta-Llama-3-8B-Instruct",
+        enable_prefix_caching=True,
+        gpu_memory_utilization=0.90,
+        enforce_eager=False,        # False means slower starts but faster inference
+        disable_log_stats=True,     # disable logging so we can stream tokens
+        disable_log_requests=True,
+    )
+
+    engine = AsyncLLMEngine.from_engine_args(engine_args)
+    duration_s = (time.monotonic_ns() - start) / 1e9
+    print(f"🏎️ engine started in {duration_s:.0f}s")
+
+    request_id = random_uuid()
+    result_generator = engine.generate(
+        prompt,
+        sampling_params,
+        request_id,
+    )
+    index, num_tokens = 0, 0
+    start = time.monotonic_ns()
+    async for output in result_generator:
+        if (
+            output.outputs[0].text
+            and "\ufffd" == output.outputs[0].text[-1]
+        ):
+            continue
+        text_delta = output.outputs[0].text[index:]
+        index = len(output.outputs[0].text)
+        num_tokens = len(output.outputs[0].token_ids)
+
+        print(text_delta)
+    duration_s = (time.monotonic_ns() - start) / 1e9
+
+    print(
+        f"\n\tGenerated {num_tokens} tokens in {duration_s:.1f}s,"
+        f" throughput = {num_tokens / duration_s:.0f} tokens/second.\n"
+    )
+
+    return
+
+
+async def xmain():
+    llm = LLM(
+        model="meta-llama/Meta-Llama-3-8B-Instruct",
+        enable_prefix_caching=True
+    )
+
+    outputs = llm.generate(prompt, sampling_params)
+
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text
+        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+
+    outputs = llm.generate(prompt, sampling_params)
+
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text
+        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())