chore: import upstream snapshot with attribution

2026-07-13 13:23:58 +08:00
commit 770d92cb1f
694 changed files with 114634 additions and 0 deletions
@@ -0,0 +1,84 @@
+import asyncio
+from typing import List  # noqa: UP035
+
+from mlc_llm.protocol.generation_config import GenerationConfig
+from mlc_llm.serve import AsyncMLCEngine, EngineConfig
+from mlc_llm.testing import require_test_model
+
+prompts = [
+    "What is the meaning of life?",
+    "Introduce the history of Pittsburgh to me. Please elaborate in detail.",
+    "Write a three-day Seattle travel plan. Please elaborate in detail.",
+    "What is Alaska famous of? Please elaborate in detail.",
+    "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.",  # noqa: E501
+    "What are the necessary components to assemble a desktop computer? Please elaborate in detail.",
+    "Why is Vitamin D important to human beings? Please elaborate in detail.",
+    "Where is milk tea originated from? Please elaborate in detail.",
+    "Where is the southernmost place in United States? Please elaborate in detail.",
+    "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.",  # noqa: E501
+]
+
+
+@require_test_model(
+    "Llama-2-7b-chat-hf-q0f16-MLC",
+    "Llama-2-7b-chat-hf-q4f16_1-MLC",
+)
+async def test_engine_generate(model: str, small_model: str):
+    # Create engine
+    async_engine = AsyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            additional_models=[small_model],
+            speculative_mode="small_draft",
+        ),
+    )
+
+    num_requests = 10
+    max_tokens = 256
+    generation_cfg = GenerationConfig(max_tokens=max_tokens)
+
+    output_texts: List[List[str]] = [  # noqa: UP006
+        ["" for _ in range(generation_cfg.n)] for _ in range(num_requests)
+    ]
+
+    async def generate_task(
+        async_engine: AsyncMLCEngine,
+        prompt: str,
+        generation_cfg: GenerationConfig,
+        request_id: str,
+    ):
+        print(f"generate task for request {request_id}")
+        rid = int(request_id)
+        async for delta_outputs in async_engine._generate(
+            prompt, generation_cfg, request_id=request_id
+        ):
+            assert len(delta_outputs) == generation_cfg.n
+            for i, delta_output in enumerate(delta_outputs):
+                output_texts[rid][i] += delta_output.delta_text
+
+    tasks = [
+        asyncio.create_task(
+            generate_task(async_engine, prompts[i], generation_cfg, request_id=str(i))
+        )
+        for i in range(num_requests)
+    ]
+
+    await asyncio.gather(*tasks)
+
+    # Print output.
+    print("All finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+    async_engine.terminate()
+    del async_engine
+
+
+if __name__ == "__main__":
+    asyncio.run(test_engine_generate())