wehub-resource-sync
206 lines
8.1 KiB
ReStructuredText
206 lines
8.1 KiB
ReStructuredText
Implement LLM Cross-engine Orchestration Patterns
|
|
======================================================================
|
|
|
|
In this tutorial, we will introduce how to implement LLM cross-engine
|
|
orchestration patterns, like prefill-decode disaggregation, in MLC-LLM
|
|
via microserving API. Aiming to make disaggregated serving programmable,
|
|
MicroServing provides a new RISC-style approach to design LLM serving
|
|
API at sub-request level. It enables programmable cross-engine serving
|
|
patterns in a few lines of python code. For more information of
|
|
microserving API, check out
|
|
https://blog.mlc.ai/2025/01/07/microserving-llm-engines.
|
|
|
|
Below is an example of prefill-decode disaggregation implementation. An
|
|
LLM cross-engine orchestration pattern is implemented in a router, which
|
|
dispatches original OpenAI-style completion requests to a chain of
|
|
microserving API calls. In this code example, we create a subclass of
|
|
Router (which includes wrappers for calling microserving APIs), and
|
|
override ``translate_request`` function. The ``translate_request``
|
|
function takes in a request and a unique identifier of the request
|
|
(``request_id``), and returns an AsyncGenerator of response. We launch
|
|
the CustomRouter and 2 engines, each of which has tensor parallel degree
|
|
2. Engine 0 is prefill engine and engine 1 is decode engine.
|
|
|
|
.. code:: python
|
|
|
|
from mlc_llm.router import Router
|
|
from mlc_llm.protocol import openai_api_protocol
|
|
from typing import Any, AsyncGenerator
|
|
from mlc_llm.serve.entrypoints import microserving_entrypoints
|
|
from mlc_llm.interface.router import serve
|
|
|
|
import aiohttp
|
|
|
|
class CustomRouter(Router):
|
|
async def translate_request(self, request: openai_api_protocol.CompletionRequest, request_id: str) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]:
|
|
pass
|
|
|
|
|
|
serve(
|
|
model="/path/to/model", # replace this with actual path
|
|
model_lib="/path/to/model_lib", # replace this with actual path
|
|
router_host="127.0.0.1",
|
|
router_port=9123,
|
|
endpoint_hosts=["127.0.0.1", "127.0.0.1"],
|
|
endpoint_ports=[9124,9125],
|
|
endpoint_num_gpus=[2,2],
|
|
enable_prefix_cache=False,
|
|
router_type=CustomRouter,
|
|
)
|
|
|
|
In the ``translate_request`` function, we first assign ``request_id`` to
|
|
request.user, and later the request id will be passed as an argument to
|
|
the microserving API.
|
|
|
|
.. code:: python
|
|
|
|
# we will pass request_id as an argument in microserving API calls
|
|
request.user = request_id
|
|
|
|
|
|
Next, call ``prep_recv`` on the decode engine to prepare KV entries for
|
|
receiving from remote. ``end=-1`` means that we will let the prefill
|
|
engine prefill all except the last token, which makes sure that the
|
|
prefill engine does not need sampling logic. ``prep_recv`` returns
|
|
address to receive KV from remote and matched prefix length. For
|
|
simplicity, we do not enable prefix cache in the tutorial, so we only
|
|
need the kv address here.
|
|
|
|
.. code:: python
|
|
|
|
async with aiohttp.ClientSession(
|
|
timeout=aiohttp.ClientTimeout(total=3 * 3600), trust_env=True
|
|
) as session:
|
|
decode_start = len(request.prompt) -1
|
|
# 1. Ask decode engine to prepare KV entries to receive from prefill engine
|
|
prep_recv_request = microserving_entrypoints.PrepRecvRequest(
|
|
**request.model_dump(), end=decode_start
|
|
)
|
|
(
|
|
kv_addr_info,
|
|
_,
|
|
) = await self.send_prepare_receive(
|
|
session=session,
|
|
request=prep_recv_request,
|
|
server_url=self.server_urls[1], # engine 0 is prefill, engine 1 is decode. Here is decode engine
|
|
)
|
|
|
|
Then, call ``remote_send`` on the prefill engine to compute and send KV
|
|
to decode engine. ``recv_rank=self.device_id_starts[1]`` means that we
|
|
are sending KV to engine 1 (decode engine).
|
|
|
|
.. code:: python
|
|
|
|
|
|
# 2. Ask prefill engine to send KV to decode engine
|
|
remote_send_request = microserving_entrypoints.RemoteSendRequest(
|
|
**request.model_dump(),
|
|
begin=0,
|
|
end=decode_start,
|
|
kv_addr_info=kv_addr_info,
|
|
recv_rank=self.device_id_starts[1], # the rank of decode engine
|
|
)
|
|
await self.send_remote_send(
|
|
session=session,
|
|
request=remote_send_request,
|
|
server_url=self.server_urls[0], # prefill engine
|
|
)
|
|
|
|
Finally, call ``start_generate`` on the decode engine to start
|
|
generating tokens. ``begin=decode_start`` means we will prefill the last
|
|
token in the prompt and start decoding. Notably, the decode process of
|
|
the request may be preempted. In such case, we yield None, so that the
|
|
router will rerun the ``translate_request`` function.
|
|
|
|
.. code:: python
|
|
|
|
# 3. Start decoding
|
|
start_generate_request = microserving_entrypoints.StartGenerateRequest(
|
|
**request.model_dump(),
|
|
begin=decode_start,
|
|
)
|
|
async for response in self.send_start_generate(
|
|
session=session,
|
|
request=start_generate_request,
|
|
server_url=self.server_urls[1],
|
|
):
|
|
if len(response.choices) > 0:
|
|
finish_reason = response.choices[0].finish_reason
|
|
if finish_reason == "preempt":
|
|
yield None
|
|
yield response
|
|
|
|
Bringing everything together, the complete code is as below:
|
|
|
|
.. code:: python
|
|
|
|
from mlc_llm.router import Router
|
|
from mlc_llm.protocol import openai_api_protocol
|
|
from typing import Any, AsyncGenerator
|
|
from mlc_llm.serve.entrypoints import microserving_entrypoints
|
|
from mlc_llm.interface.router import serve
|
|
|
|
import aiohttp
|
|
class CustomRouter(Router):
|
|
async def translate_request(self, request: openai_api_protocol.CompletionRequest, request_id: str) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]:
|
|
# we will pass request_id as an argument in microserving API calls
|
|
request.user = request_id
|
|
|
|
async with aiohttp.ClientSession(
|
|
timeout=aiohttp.ClientTimeout(total=3 * 3600), trust_env=True
|
|
) as session:
|
|
decode_start = len(request.prompt) -1
|
|
# 1. Ask decode engine to prepare KV entries to receive from prefill engine
|
|
prep_recv_request = microserving_entrypoints.PrepRecvRequest(
|
|
**request.model_dump(), end=decode_start
|
|
)
|
|
(
|
|
kv_addr_info,
|
|
_,
|
|
) = await self.send_prepare_receive(
|
|
session=session,
|
|
request=prep_recv_request,
|
|
server_url=self.server_urls[1], # engine 0 is prefill, engine 1 is decode. Here is decode engine
|
|
)
|
|
# 2. Ask prefill engine to send KV to decode engine
|
|
remote_send_request = microserving_entrypoints.RemoteSendRequest(
|
|
**request.model_dump(),
|
|
begin=0,
|
|
end=decode_start,
|
|
kv_addr_info=kv_addr_info,
|
|
recv_rank=self.device_id_starts[1], # the rank of decode engine
|
|
)
|
|
await self.send_remote_send(
|
|
session=session,
|
|
request=remote_send_request,
|
|
server_url=self.server_urls[0], # prefill engine
|
|
)
|
|
# 3. Start decoding
|
|
start_generate_request = microserving_entrypoints.StartGenerateRequest(
|
|
**request.model_dump(),
|
|
begin=decode_start,
|
|
)
|
|
async for response in self.send_start_generate(
|
|
session=session,
|
|
request=start_generate_request,
|
|
server_url=self.server_urls[1],
|
|
):
|
|
if len(response.choices) > 0:
|
|
finish_reason = response.choices[0].finish_reason
|
|
if finish_reason == "preempt":
|
|
yield None
|
|
yield response
|
|
|
|
|
|
serve(
|
|
model="/path/to/model", # replace this with actual path
|
|
model_lib="/path/to/model_lib", # replace this with actual path
|
|
router_host="127.0.0.1",
|
|
router_port=9123,
|
|
endpoint_hosts=["127.0.0.1", "127.0.0.1"],
|
|
endpoint_ports=[9124,9125],
|
|
endpoint_num_gpus=[2,2],
|
|
enable_prefix_cache=False,
|
|
router_type=CustomRouter,
|
|
)
|