wehub/mlc-ai--mlc-llm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
Build Docs / Deploy Docs (push) Has been cancelled
Details
Windows CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:23:58 +08:00
commit 770d92cb1f
694 changed files with 114634 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mlc_llm/serve/sync_engine.py
+363
View File
@@ -0,0 +1,363 @@
"""The MLC LLM synchronized engine.
NOTE: This engine defined in this file directly wraps the underlying
Engine implementation in C++, is not optimized by multi-threading and
does not offer standard OpenAI API interface.
We do not expose it and use it by default. As of now it mainly serves
the test and debug purpose because of its simplicity.
"""
import json
from collections.abc import Sequence
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union # noqa: UP035
import tvm
from mlc_llm.protocol.generation_config import GenerationConfig
from mlc_llm.serve import data
from mlc_llm.serve.config import EngineConfig
from mlc_llm.serve.engine_base import (
EngineMetrics,
_check_engine_config,
_parse_models,
_print_engine_mode_logging_msg,
_process_model_args,
detect_device,
)
from mlc_llm.serve.event_trace_recorder import EventTraceRecorder
from mlc_llm.serve.request import Request
from mlc_llm.support import logging
from mlc_llm.tokenizers import TextStreamer, Tokenizer
logger = logging.getLogger(__name__)
def _create_tvm_module(
creator: str,
ffi_funcs: Sequence[str],
creator_args: Optional[List[Any]] = None, # noqa: UP006
) -> Dict[str, Callable]: # noqa: UP006
"""Internal method to create a module."""
if creator_args is None:
creator_args = []
module = tvm.get_global_func(creator, allow_missing=False)(*creator_args)
return {key: module[key] for key in ffi_funcs}
class SyncMLCEngine:
"""The Python interface of synchronize request serving engine for MLC LLM.
The engine receives requests from the "add_request" method. For
an given request, the engine will keep generating new tokens for
the request until finish (under certain criterion). After finish,
the engine will return the generation result through the callback
function provided by the request.
NOTE: This engine directly wraps the underlying Engine implementation
in C++, is not optimized by multi-threading and does not offer standard
OpenAI API interface. We do not expose it and use it by default.
As of now it mainly serves the test and debug purpose because of its
simplicity.
Parameters
----------
engine_config : Optional[EngineConfig]
Additional configurable arguments of MLC engine.
See class "EngineConfig" for more detail.
enable_tracing : bool
A boolean indicating if to enable event logging for requests.
request_stream_callback : Optional[Callable[[str, data.TokenData, Optional[str]], None]]
The provided callback function to handle the generation
output. It has the signature of `(str, data.TokenData, bool) -> None`,
where
- the first string is the request id,
- the TokenData contains the generated **delta** token ids since
the last invocation of the callback on the specific request,
- the optional string value denotes the finish reason if the
generation of the request is finished, or None if it has not finished.
The callback function is optional at construction, but it needs to
be set before the engine executing requests. This can be done via
the `set_request_stream_callback` method. Otherwise, the engine will raise
exception.
"""
def __init__(
self,
model: str,
device: Union[str, tvm.runtime.Device] = "auto",
*,
model_lib: Optional[str] = None,
mode: Literal["local", "interactive", "server"] = "local",
engine_config: Optional[EngineConfig] = None,
enable_tracing: bool = False,
request_stream_callback: Optional[Callable[[List[data.RequestStreamOutput]], None]] = None, # noqa: UP006
):
# - Check the fields fields of `engine_config`.
if engine_config is None:
engine_config = EngineConfig()
_check_engine_config(
model,
model_lib,
mode,
engine_config,
)
# - Initialize model loading info.
models = _parse_models(model, model_lib, engine_config.additional_models)
if isinstance(device, str):
device = detect_device(device)
assert isinstance(device, tvm.runtime.Device)
(
model_args,
model_config_paths,
self.conv_template,
) = _process_model_args(models, device, engine_config)
# - Load the raw model config into dict
self.model_config_dicts = []
for i, model_info in enumerate(models):
model_info.model_lib = model_args[i][1]
with open(model_config_paths[i], encoding="utf-8") as file:
self.model_config_dicts.append(json.load(file))
# - Print logging info for regarding the mode selection.
if engine_config.verbose:
_print_engine_mode_logging_msg(mode)
self._ffi = _create_tvm_module(
"mlc.serve.create_engine",
ffi_funcs=[
"init",
"add_request",
"abort_request",
"step",
"reset",
"json_metrics",
"get_request_stream_callback",
"set_request_stream_callback",
"create_request",
],
)
self.trace_recorder = EventTraceRecorder() if enable_tracing else None
engine_config.model = model_args[0][0]
engine_config.model_lib = model_args[0][1]
engine_config.additional_models = model_args[1:]
engine_config.mode = mode
self._ffi["init"](
engine_config.asjson(),
device,
request_stream_callback,
self.trace_recorder,
)
self.tokenizer = Tokenizer(model_args[0][0])
def generate(
self,
prompts: Union[str, List[str], List[int], List[List[int]], List[List[data.Data]]], # noqa: UP006
generation_config: Union[GenerationConfig, List[GenerationConfig]], # noqa: UP006
) -> Tuple[List[List[str]], List[Optional[List[List[str]]]]]: # noqa: UP006
"""Generate texts for a list of input prompts.
Each prompt can be a string or a list of token ids.
The generation for each prompt is independent.
Return the generation results, one for each prompt.
Parameters
----------
prompts : Union[str, List[str], List[int], List[List[int]]]
One or a list of input prompts for text generation.
Each prompt can be a string or a list of token ids.
generation_config : Union[GenerationConfig, List[GenerationConfig]]
The generation config for each requests.
If the it is a single GenerationConfig instance,
this config will be shared by all the prompts.
Otherwise, one generation config is required for every
prompt.
Returns
-------
output_text : List[List[str]]
The text generation results, one list of strings for each input prompt.
The length of each list is the parallel generation `n` in
generation config.
output_logprobs_str : List[Optional[List[List[str]]]]
The logprob strings of each token for each input prompt, or None
if an input prompt does not require logprobs.
"""
if isinstance(prompts, str):
# `prompts` is a single string.
prompts = [prompts]
else:
assert isinstance(prompts, list), (
"Input `prompts` is expected to be a string, a list of "
"str, a list of token ids or multiple lists of token ids. "
)
if len(prompts) == 0:
return [], []
if isinstance(prompts[0], int):
# `prompts` is a list of token ids
prompts = [prompts]
num_requests = len(prompts)
if not isinstance(generation_config, list):
generation_config = [generation_config] * num_requests
assert len(generation_config) == num_requests, (
"Number of generation config and number of prompts mismatch"
)
num_finished_generations = 0
output_texts: List[List[str]] = [] # noqa: UP006
output_logprobs_str: List[Optional[List[List[str]]]] = [] # noqa: UP006
text_streamers: List[List[TextStreamer]] = [] # noqa: UP006
for i in range(num_requests):
output_texts.append([])
output_logprobs_str.append([] if generation_config[i].logprobs else None)
text_streamers.append([])
for _ in range(generation_config[i].n):
output_texts[i].append("")
text_streamers[i].append(TextStreamer(self.tokenizer))
if output_logprobs_str[i] is not None:
output_logprobs_str[i].append([])
num_total_generations = sum(cfg.n for cfg in generation_config)
# Save a copy of the original function callback since `generate`
# overrides the callback function.
# The original callback will be set back later on.
original_callback = self._ffi["get_request_stream_callback"]()
# Define the callback function for request generation results
def request_stream_callback(delta_outputs: List[data.RequestStreamOutput]): # noqa: UP006
nonlocal num_finished_generations
for delta_output in delta_outputs:
request_id, stream_outputs = delta_output.unpack()
rid = int(request_id)
assert len(stream_outputs) == generation_config[rid].n
for i, (stream_output, text_streamer) in enumerate(
zip(stream_outputs, text_streamers[rid])
):
if output_logprobs_str[rid] is not None:
assert stream_output.delta_logprob_json_strs is not None
output_logprobs_str[rid][i] += stream_output.delta_logprob_json_strs
delta_text = stream_output.extra_prefix_string + (
text_streamer.put(stream_output.delta_token_ids)
if len(stream_output.delta_token_ids) > 0
else ""
)
if stream_output.finish_reason is not None:
delta_text += text_streamer.finish()
output_texts[rid][i] += delta_text
if stream_output.finish_reason is not None:
num_finished_generations += 1
# Override the callback function in engine.
self._ffi["set_request_stream_callback"](request_stream_callback)
def convert_to_data(
prompt: Union[str, List[int], List[data.Data]], # noqa: UP006
) -> List[data.Data]: # noqa: UP006
if isinstance(prompt, str):
return [data.TextData(prompt)]
if isinstance(prompt[0], int):
return [data.TokenData(prompt)]
return prompt
# Add requests to engine.
for req_id, (prompt, generation_cfg) in enumerate(zip(prompts, generation_config)):
input_data = convert_to_data(prompt)
self.add_request(
self.create_request(
request_id=str(req_id),
inputs=input_data,
generation_config=generation_cfg,
)
)
while num_finished_generations != num_total_generations:
self.step()
# Restore the callback function in engine.
self._ffi["set_request_stream_callback"](original_callback)
return output_texts, output_logprobs_str
def create_request(
self,
request_id: str,
inputs: Union[data.Data, List[data.Data]], # noqa: UP006
generation_config: GenerationConfig,
):
"""Create a new request that can be added to engine.
Parameters
----------
request_id : str
The unique identifier of the request.
Different requests should have different ids.
inputs : List[Data]
The user inputs of a request. Input may have multi-modality.
generation_config : GenerationConfig
The generation configuration of the request.
Note
----
engine may fill in default generation config of the model.
"""
if not isinstance(inputs, list):
inputs = [inputs]
return self._ffi["create_request"](
request_id, inputs, generation_config.model_dump_json(by_alias=True)
)
def add_request(self, request: Request) -> None:
"""Add a new request to the engine.
Parameters
----------
request : Request
The request to add.
"""
self._ffi["add_request"](request)
def abort_request(self, request_id: str) -> None:
"""Abort the generation of the request corresponding to the input request id.
Parameters
----------
request_id : str
The unique id of the request to abort.
"""
self._ffi["abort_request"](request_id)
def step(self) -> None:
"""The main function that the engine takes a step of action.
At each step, the engine may decide to
- run prefill for one (or more) requests,
- run one-step decode for the all existing requests
...
In the end of certain actions (e.g., decode), the engine will
check if any request has finished, and will return the
generation results for those finished requests.
"""
self._ffi["step"]()
def reset(self) -> None:
"""Reset the engine, clean up all running data and metrics."""
self._ffi["reset"]()
def metrics(self) -> EngineMetrics:
"""Reset the engine, clean up all running data and metrics."""
return EngineMetrics(json.loads(self._ffi["json_metrics"]()))