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,101 @@
+import asyncio
+
+from mlc_llm.protocol import openai_api_protocol
+from mlc_llm.router import Router
+
+model_tp1 = "./dist/Llama-3.2-1B-Instruct-q0f16-MLC/"
+model_lib_tp1 = "./dist/lib/Llama-3.2-1B-q0f16-cuda.so"
+# model_lib_tp1 = None
+
+model_tp2 = "./dist/Llama-3.2-1B-Instruct-q0f16-MLC-tp2/"
+model_lib_tp2 = "./dist/lib/Llama-3.2-1B-q0f16-cuda-tp2.so"
+# model_lib_tp2 = None
+
+
+def get_router_1tp1():
+    return (
+        Router(
+            model_tp1,
+            model_lib=model_lib_tp1,
+            hosts=["127.0.0.1"],
+            ports=[8080],
+        ),
+        model_tp1,
+    )
+
+
+def get_router_2tp1():
+    return (
+        Router(
+            model_tp1,
+            model_lib=model_lib_tp1,
+            hosts=["127.0.0.1", "127.0.0.1"],
+            ports=[8080, 8081],
+            device_id_starts=[0, 1],
+            npes=2,
+        ),
+        model_tp1,
+    )
+
+
+def get_router_1tp2():
+    return (
+        Router(
+            model_tp2,
+            model_lib=model_lib_tp2,
+            hosts=["127.0.0.1"],
+            ports=[8080],
+            npes=2,
+        ),
+        model_tp2,
+    )
+
+
+def get_router_2tp2():
+    return (
+        Router(
+            model_tp2,
+            model_lib=model_lib_tp2,
+            hosts=["127.0.0.1", "127.0.0.1"],
+            ports=[8080, 8081],
+            device_id_starts=[0, 2],
+            npes=4,
+        ),
+        model_tp2,
+    )
+
+
+CONFIG_TO_ROUTER = {
+    "1tp1": get_router_1tp1,
+    "2tp1": get_router_2tp1,
+    "1tp2": get_router_1tp2,
+    "2tp2": get_router_2tp2,
+}
+
+
+async def test_router(schedule: str = "round_robin", endpoints_config: str = "1tp1"):
+    router, model_id = CONFIG_TO_ROUTER[endpoints_config]()
+
+    request = openai_api_protocol.CompletionRequest(
+        prompt="The meaning of life ",
+        model=model_id,
+        stream=True,
+        max_tokens=64,
+        stream_options=openai_api_protocol.StreamOptions(include_usage=True),
+    )
+    if schedule == "round_robin":
+        async for chunk in router._handle_completion_round_robin(request, "1"):
+            print(chunk)
+    elif schedule == "disagg":
+        async for chunk in router._handle_completion_disagg(request, "1"):
+            print(chunk)
+    else:
+        raise ValueError(f"Unknown scheduling method: {schedule}")
+    router.terminate()
+
+
+if __name__ == "__main__":
+    # asyncio.run(test_router("round_robin", endpoints_config="1tp1"))
+    # asyncio.run(test_router("round_robin", endpoints_config="1tp2"))
+    # asyncio.run(test_router("round_robin", endpoints_config="2tp1"))
+    asyncio.run(test_router("round_robin", endpoints_config="2tp2"))