chore: import upstream snapshot with attribution

python/mlc_llm/serve/entrypoints/__init__.py

@@ -0,0 +1,8 @@
+"""The entrypoints for MLC LLM server."""
+
+from . import (
+    debug_entrypoints,
+    metrics_entrypoints,
+    microserving_entrypoints,
+    openai_entrypoints,
+)

python/mlc_llm/serve/entrypoints/debug_entrypoints.py

@@ -0,0 +1,128 @@
+"""MLC LLM server debug entrypoints"""
+
+import json
+from http import HTTPStatus
+
+import fastapi
+
+from mlc_llm.protocol import error_protocol
+from mlc_llm.serve.server import ServerContext
+
+app = fastapi.APIRouter()
+
+################ /debug/dump_event_trace ################
+
+
+@app.post("/debug/dump_event_trace")
+async def debug_dump_event_trace(request: fastapi.Request):
+    """Return the recorded events in Chrome Trace Event Format in JSON string.
+    The input request payload should have only one field, specifying the
+    model to query. For example: `{"model": "Llama-2-7b-chat-hf-q0f16"}`.
+    """
+    # Get the raw request body as bytes
+    request_raw_data = await request.body()
+    request_json_str = request_raw_data.decode("utf-8")
+    try:
+        # Parse the JSON string
+        request_dict = json.loads(request_json_str)
+    except json.JSONDecodeError:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
+        )
+    if "model" not in request_dict:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
+        )
+
+    # Check the requested model.
+    model = request_dict["model"]
+
+    server_context: ServerContext = ServerContext.current()
+    async_engine = server_context.get_engine(model)
+
+    if async_engine is None:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            message=f'The requested model "{model}" is not served.',
+        )
+    if async_engine.state.trace_recorder is None:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            message=f'The requested model "{model}" does not enable tracing',
+        )
+
+    return json.loads(async_engine.state.trace_recorder.dump_json())
+
+
+################ /debug/cuda_profiler_start/end ################
+
+
+@app.post("/debug/cuda_profiler_start")
+async def debug_cuda_profiler_start(_request: fastapi.Request):
+    """Start the cuda profiler for the engine. Only for debug purpose."""
+    server_context: ServerContext = ServerContext.current()
+    # Since the CUDA profiler is process-wise, call the function for one model is sufficient.
+    for model in server_context.get_model_list():
+        async_engine = server_context.get_engine(model)
+        async_engine._debug_call_func_on_all_worker("mlc.debug_cuda_profiler_start")
+        break
+
+
+@app.post("/debug/cuda_profiler_stop")
+async def debug_cuda_profiler_stop(_request: fastapi.Request):
+    """Stop the cuda profiler for the engine. Only for debug purpose."""
+    server_context: ServerContext = ServerContext.current()
+    # Since the CUDA profiler is process-wise, call the function for one model is sufficient.
+    for model in server_context.get_model_list():
+        async_engine = server_context.get_engine(model)
+        async_engine._debug_call_func_on_all_worker("mlc.debug_cuda_profiler_stop")
+        break
+
+
+@app.post("/debug/dump_engine_metrics")
+async def debug_dump_engine_metrics(request: fastapi.Request):
+    """Dump the engine metrics for the engine. Only for debug purpose."""
+    # Get the raw request body as bytes
+    request_raw_data = await request.body()
+    request_json_str = request_raw_data.decode("utf-8")
+    try:
+        # Parse the JSON string
+        request_dict = json.loads(request_json_str)
+    except json.JSONDecodeError:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
+        )
+
+    # Check the requested model.
+    model = request_dict.get("model", None)
+
+    server_context: ServerContext = ServerContext.current()
+    async_engine = server_context.get_engine(model)
+    res = await async_engine.metrics()
+    return res
+
+
+@app.post("/debug/reset_engine")
+async def debug_reset_engine_stats(request: fastapi.Request):
+    """Reset the engine, clean up all running data and metrics."""
+    # Get the raw request body as bytes
+    request_raw_data = await request.body()
+    request_json_str = request_raw_data.decode("utf-8")
+    try:
+        # Parse the JSON string
+        request_dict = json.loads(request_json_str)
+    except json.JSONDecodeError:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
+        )
+    if "model" not in request_dict:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
+        )
+
+    # Check the requested model.
+    model = request_dict["model"]
+
+    server_context: ServerContext = ServerContext.current()
+    async_engine = server_context.get_engine(model)
+    async_engine.reset()

python/mlc_llm/serve/entrypoints/metrics_entrypoints.py

@@ -0,0 +1,23 @@
+"""MLC LLM server metrics entrypoints"""
+
+import fastapi
+from fastapi.responses import PlainTextResponse
+
+from mlc_llm.serve.server import ServerContext
+
+app = fastapi.APIRouter()
+
+################ /metrics ################
+
+
+@app.get("/metrics", response_class=PlainTextResponse)
+async def metrics(_request: fastapi.Request):
+    """Start the cuda profiler for the engine. Only for debug purpose."""
+    server_context: ServerContext = ServerContext.current()
+    # Use the metrics from first engine for now
+    # TODO(mlc-team): consider refactor server context to
+    # single engine since multiple AsyncMLCEngine do not work well with each other
+    # We need to work within the internal engine instead.
+    for model in server_context.get_model_list():
+        async_engine = server_context.get_engine(model)
+        return (await async_engine.metrics()).prometheus_text()

python/mlc_llm/serve/entrypoints/microserving_entrypoints.py

@@ -0,0 +1,73 @@
+"""MicroServing server entrypoints in MLC LLM"""
+
+import fastapi
+
+from mlc_llm.protocol.debug_protocol import DisaggConfig
+from mlc_llm.protocol.microserving_protocol import (
+    PrepRecvRequest,
+    PrepRecvResponse,
+    RemoteSendRequest,
+    StartGenerateRequest,
+)
+from mlc_llm.protocol.openai_api_protocol import StreamOptions
+
+from .openai_entrypoints import request_completion
+
+app = fastapi.APIRouter()
+
+
+################ MicroServing Endpoints ################
+
+
+@app.post("/microserving/prep_recv")
+async def prep_recv(request: PrepRecvRequest, raw_request: fastapi.Request) -> PrepRecvResponse:
+    """Handle the microserving request for receive preparation.
+    Match the prompt in the prefix cache (when enabled),
+    allocate entries in the KV cache to prepare receiving the KV data of the prompt.
+    Return the matched prefix length and the allocated KV entry metadata.
+    """
+    request.debug_config.disagg_config = DisaggConfig(
+        kind="prepare_receive",
+        kv_window_begin=0,  # always zero for prepare_receive
+        kv_window_end=request.end,
+    )
+    request.stream_options = StreamOptions(include_usage=True)
+    request.stream = False
+
+    response = await request_completion(request=request, raw_request=raw_request)
+    assert response.usage is not None
+    assert response.usage.extra is not None
+    assert "prefix_matched_length" in response.usage.extra
+    assert "kv_append_metadata" in response.usage.extra
+    return PrepRecvResponse(
+        prefix_matched_length=response.usage.extra["prefix_matched_length"],
+        kv_append_metadata=response.usage.extra["kv_append_metadata"],
+    )
+
+
+@app.post("/microserving/remote_send")
+async def remote_send(request: RemoteSendRequest, raw_request: fastapi.Request):
+    """Compute and generate the KV data of the prompt in the specified KV window.
+    Send the KV data to the destination server."""
+    request.debug_config.disagg_config = DisaggConfig(
+        kind="remote_send",
+        kv_window_begin=request.begin,
+        kv_window_end=request.end,
+        kv_append_metadata=request.kv_addr_info,
+        dst_group_offset=request.recv_rank,
+    )
+    request.stream_options = StreamOptions(include_usage=True)
+    request.stream = False
+
+    await request_completion(request=request, raw_request=raw_request)
+    return {}
+
+
+@app.post("/microserving/start_generate")
+async def start_generate(request: StartGenerateRequest, raw_request: fastapi.Request):
+    """Prefill the prompt in the specified KV window, and start decode."""
+    request.debug_config.disagg_config = DisaggConfig(
+        kind="start_generation",
+        kv_window_begin=request.begin,
+    )
+    return await request_completion(request=request, raw_request=raw_request)

python/mlc_llm/serve/entrypoints/openai_entrypoints.py

@@ -0,0 +1,340 @@
+"""OpenAI API-compatible server entrypoints in MLC LLM"""
+
+import base64
+import struct
+from collections.abc import AsyncGenerator
+from http import HTTPStatus
+from typing import List, Optional  # noqa: UP035
+
+import fastapi
+import numpy as np
+
+from mlc_llm.protocol import error_protocol
+from mlc_llm.protocol.openai_api_protocol import (
+    ChatCompletionRequest,
+    CompletionLogProbs,
+    CompletionRequest,
+    EmbeddingObject,
+    EmbeddingRequest,
+    EmbeddingResponse,
+    EmbeddingUsage,
+    ListResponse,
+    LogProbsContent,
+    ModelResponse,
+)
+from mlc_llm.serve import engine_base, engine_utils
+from mlc_llm.serve.server import ServerContext
+
+
+def verify_api_key(request: fastapi.Request):
+    """Function to verify API key"""
+    server_context = ServerContext.current()
+    # Only perform verification when API key is configured
+    if server_context is not None and server_context.api_key is not None:
+        provided_key = request.headers.get("Authorization", "").replace("Bearer ", "")
+        if provided_key != server_context.api_key:
+            raise fastapi.HTTPException(status_code=401, detail="Invalid API Key")
+
+
+app = fastapi.APIRouter(dependencies=[fastapi.Depends(verify_api_key)])
+
+
+################ v1/embeddings ################
+
+
+@app.post("/v1/embeddings")
+async def request_embedding(request: EmbeddingRequest):
+    """OpenAI-compatible embedding API.
+    API reference: https://platform.openai.com/docs/api-reference/embeddings/create
+    """
+    server_context: ServerContext = ServerContext.current()
+    embedding_engine = server_context.get_embedding_engine(request.model)
+    if embedding_engine is None:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            message=f'The requested model "{request.model}" is not served as an embedding model.',
+        )
+
+    # Normalize input to List[str]
+    inputs: List[str]  # noqa: UP006
+    if isinstance(request.input, str):
+        inputs = [request.input]
+    elif (
+        isinstance(request.input, list)
+        and len(request.input) > 0
+        and isinstance(request.input[0], str)
+    ):
+        inputs = list(request.input)
+    else:
+        # Token ID inputs (List[int] or List[List[int]]) — decode back to strings
+        if isinstance(request.input[0], int):
+            inputs = [embedding_engine.tokenizer.decode(request.input)]
+        else:
+            inputs = [embedding_engine.tokenizer.decode(ids) for ids in request.input]
+
+    # Run embedding inference (async — does not block the event loop)
+    try:
+        embeddings, total_tokens = await embedding_engine.async_embed(inputs)
+    except Exception as exc:
+        return error_protocol.create_error_response(
+            HTTPStatus.INTERNAL_SERVER_ERROR,
+            message=f"Embedding inference failed: {exc}",
+        )
+
+    # Optional: truncate dimensions (Matryoshka-style).
+    # This is API-level renormalization after dimension truncation,
+    # independent of model metadata normalize. Always renormalize
+    # truncated vectors to maintain unit length per OpenAI API contract.
+    if request.dimensions is not None:
+        for i, emb in enumerate(embeddings):
+            vec = np.array(emb[: request.dimensions], dtype=np.float32)
+            norm = np.linalg.norm(vec)
+            if norm > 1e-12:
+                vec = vec / norm
+            embeddings[i] = vec.tolist()
+
+    # Build response data
+    resp_data = []
+    for i, emb in enumerate(embeddings):
+        if request.encoding_format == "base64":
+            binary = struct.pack(f"<{len(emb)}f", *emb)
+            resp_data.append(
+                EmbeddingObject(
+                    embedding=base64.b64encode(binary).decode("utf-8"),
+                    index=i,
+                )
+            )
+        else:
+            resp_data.append(EmbeddingObject(embedding=emb, index=i))
+
+    return EmbeddingResponse(
+        data=resp_data,
+        model=request.model,
+        usage=EmbeddingUsage(prompt_tokens=total_tokens, total_tokens=total_tokens),
+    )
+
+
+################ v1/models ################
+
+
+@app.get("/v1/models")
+async def request_models() -> ListResponse:
+    """OpenAI-compatible served model query API.
+    API reference: https://platform.openai.com/docs/api-reference/models
+    """
+    server_context: ServerContext = ServerContext.current()
+    return ListResponse(data=[ModelResponse(id=model) for model in server_context.get_model_list()])
+
+
+################ v1/completions ################
+
+
+@app.post("/v1/completions")
+async def request_completion(request: CompletionRequest, raw_request: fastapi.Request):
+    """OpenAI-compatible completion API.
+    API reference: https://platform.openai.com/docs/api-reference/completions/create
+    """
+    # - Check the requested model.
+    server_context: ServerContext = ServerContext.current()
+    request_final_usage_include_extra = server_context.enable_debug
+    request_include_debug_config = server_context.enable_debug
+
+    if not request_include_debug_config:
+        request.debug_config = None
+
+    async_engine = server_context.get_engine(request.model)
+    if async_engine is None:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            message=f'The requested model "{request.model}" is not served.',
+        )
+    # FIXME: This is a temporary solution to make sure
+    # prep_recv, remote_send and start_generation process the same request
+    request_id = request.user if request.user is not None else f"cmpl-{engine_utils.random_uuid()}"
+
+    # Streaming response.
+    if request.stream:
+        # We manually get the first response from generator to
+        # capture potential exceptions in this scope, rather then
+        # the StreamingResponse scope.
+        stream_generator = async_engine._handle_completion(
+            request,
+            request_id,
+            request_final_usage_include_extra=request_final_usage_include_extra,
+        )
+        first_response = await anext(  # noqa: F821
+            stream_generator
+        )
+
+        async def completion_stream_generator() -> AsyncGenerator[str, None]:
+            if isinstance(first_response, StopAsyncIteration):
+                yield "data: [DONE]\n\n"
+                return
+            yield f"data: {first_response.model_dump_json(by_alias=True)}\n\n"
+            async for response in stream_generator:
+                yield f"data: {response.model_dump_json(by_alias=True)}\n\n"
+            yield "data: [DONE]\n\n"
+
+        return fastapi.responses.StreamingResponse(
+            completion_stream_generator(), media_type="text/event-stream"
+        )
+
+    # Normal response.
+    request_final_usage = None
+    output_texts = [""] * request.n
+    finish_reasons: List[Optional[str]] = [None] * request.n  # noqa: UP006
+    logprob_results: List[Optional[CompletionLogProbs]] = [None] * request.n  # noqa: UP006
+
+    async for response in async_engine._handle_completion(
+        request,
+        request_id,
+        request_final_usage_include_extra=request_final_usage_include_extra,
+    ):
+        if await raw_request.is_disconnected():
+            # In non-streaming cases, the engine will not be notified
+            # when the request is disconnected.
+            # Therefore, we check if it is disconnected each time,
+            # and explicitly return.
+            # Note that requesta abort is triggered when the async for and funciton scope ends.
+            return error_protocol.create_error_response(
+                HTTPStatus.BAD_REQUEST, message="The request has disconnected"
+            )
+        # this is the final chunk
+        if response.usage is not None:
+            request_final_usage = response.usage
+            # remove extra information if debug is not enabled
+            if not server_context.enable_debug:
+                request_final_usage.extra = None
+            continue
+        for choice in response.choices:
+            output_texts[choice.index] += choice.text
+            if choice.finish_reason is not None and finish_reasons[choice.index] is None:
+                finish_reasons[choice.index] = choice.finish_reason
+            if choice.logprobs is not None:
+                if logprob_results[choice.index] is None:
+                    logprob_results[choice.index] = choice.logprobs
+                else:
+                    logprob_results[choice.index].token_logprobs.extend(
+                        choice.logprobs.token_logprobs
+                    )
+                    logprob_results[choice.index].tokens.extend(choice.logprobs.tokens)
+                    logprob_results[choice.index].top_logprobs.extend(choice.logprobs.top_logprobs)
+
+    return engine_base.wrap_completion_response(
+        request_id=request_id,
+        model=request.model,
+        output_texts=output_texts,
+        finish_reasons=finish_reasons,
+        logprob_results=logprob_results,
+        usage=request_final_usage,
+    )
+
+
+################ v1/chat/completions ################
+
+
+@app.post("/v1/chat/completions")
+async def request_chat_completion(request: ChatCompletionRequest, raw_request: fastapi.Request):
+    """OpenAI-compatible chat completion API.
+    API reference: https://platform.openai.com/docs/api-reference/chat
+    """
+    # - Check the requested model.
+    server_context: ServerContext = ServerContext.current()
+    request_final_usage_include_extra = server_context.enable_debug
+    request_include_debug_config = server_context.enable_debug
+
+    if not request_include_debug_config:
+        request.debug_config = None
+
+    async_engine = server_context.get_engine(request.model)
+    if async_engine is None:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            message=f'The requested model "{request.model}" is not served.',
+        )
+    # FIXME: This is a temporary solution to make sure
+    # prep_recv, remote_send and start_generation process the same request
+    request_id = (
+        request.user if request.user is not None else f"chatcmpl-{engine_utils.random_uuid()}"
+    )
+
+    # Streaming response.
+    if request.stream:
+        # We manually get the first response from generator to
+        # capture potential exceptions in this scope, rather then
+        # the StreamingResponse scope.
+        stream_generator = async_engine._handle_chat_completion(
+            request,
+            request_id,
+            request_final_usage_include_extra=request_final_usage_include_extra,
+        )
+        first_response = await anext(  # noqa: F821
+            stream_generator
+        )
+
+        async def completion_stream_generator() -> AsyncGenerator[str, None]:
+            if isinstance(first_response, StopAsyncIteration):
+                yield "data: [DONE]\n\n"
+                return
+            yield f"data: {first_response.model_dump_json(by_alias=True)}\n\n"
+            async for response in stream_generator:
+                yield f"data: {response.model_dump_json(by_alias=True)}\n\n"
+            yield "data: [DONE]\n\n"
+
+        return fastapi.responses.StreamingResponse(
+            completion_stream_generator(), media_type="text/event-stream"
+        )
+
+    # Normal response.
+    request_final_usage = None
+    output_texts = ["" for _ in range(request.n)]
+    finish_reasons: List[Optional[str]] = [None for _ in range(request.n)]  # noqa: UP006
+    logprob_results: Optional[List[List[LogProbsContent]]] = (  # noqa: UP006
+        [[] for _ in range(request.n)] if request.logprobs else None
+    )
+
+    async for response in async_engine._handle_chat_completion(
+        request,
+        request_id,
+        request_final_usage_include_extra=request_final_usage_include_extra,
+    ):
+        if await raw_request.is_disconnected():
+            # In non-streaming cases, the engine will not be notified
+            # when the request is disconnected.
+            # Therefore, we check if it is disconnected each time,
+            # no need to explicitly abort, as the chat completion
+            # return will trigger abort call
+            return error_protocol.create_error_response(
+                HTTPStatus.BAD_REQUEST, message="The request has disconnected"
+            )
+        # usage is always the last chunk
+        if response.usage is not None:
+            request_final_usage = response.usage
+            # remove extra information if debug is not enabled
+            if not server_context.enable_debug:
+                request_final_usage.extra = None
+
+        for choice in response.choices:
+            assert isinstance(choice.delta.content, str)
+            output_texts[choice.index] += choice.delta.content
+            if choice.finish_reason is not None and finish_reasons[choice.index] is None:
+                finish_reasons[choice.index] = choice.finish_reason
+            if choice.logprobs is not None:
+                assert logprob_results is not None
+                logprob_results[choice.index] += choice.logprobs.content
+
+    assert all(finish_reason is not None for finish_reason in finish_reasons)
+    use_function_calling, tool_calls_list = engine_base.process_function_call_output(
+        output_texts, finish_reasons
+    )
+
+    return engine_base.wrap_chat_completion_response(
+        request_id=request_id,
+        model=request.model,
+        output_texts=output_texts,
+        finish_reasons=finish_reasons,
+        tool_calls_list=tool_calls_list,
+        logprob_results=logprob_results,
+        use_function_calling=use_function_calling,
+        usage=request_final_usage,
+    )