Implementation:Mlc ai Mlc llm Serve Data Py

Overview

python/mlc_llm/serve/data.py defines the multi-modality data classes used throughout the MLC LLM serving infrastructure. It provides a type hierarchy for representing text, token, and image data as TVM runtime objects, along with dataclasses for streaming request outputs. These classes serve as the core data abstractions that flow between the Python serving layer and the underlying C++ engine via TVM's Foreign Function Interface (FFI).

Location

• File: python/mlc_llm/serve/data.py
• Module: mlc_llm.serve.data
• Lines: 216

Key Components

Data (Base Class)

Data is the abstract base class for all multi-modality data types. It inherits from tvm.runtime.Object and is registered with TVM FFI under the name mlc.serve.Data.

@tvm_ffi.register_object("mlc.serve.Data")
class Data(Object):
    """The base class of multi-modality data (text, tokens, embedding, etc)."""

    def __init__(self):
        pass

TextData

TextData wraps a plain text string. It is registered as mlc.serve.TextData and delegates construction and property access to C++ FFI functions.

@tvm_ffi.register_object("mlc.serve.TextData")
class TextData(Data):
    def __init__(self, text: str):
        self.__init_handle_by_constructor__(_ffi_api.TextData, text)

    @property
    def text(self) -> str:
        return str(_ffi_api.TextDataGetTextString(self))

• Constructor: Delegates to _ffi_api.TextData on the C++ side.
• text property: Returns the underlying text string via _ffi_api.TextDataGetTextString.
• __str__: Returns the text value directly.

TokenData

TokenData holds a list of token IDs (integers). It is registered as mlc.serve.TokenData.

@tvm_ffi.register_object("mlc.serve.TokenData")
class TokenData(Data):
    def __init__(self, token_ids: List[int]):
        self.__init_handle_by_constructor__(_ffi_api.TokenData, *token_ids)

    @property
    def token_ids(self) -> List[int]:
        return list(_ffi_api.TokenDataGetTokenIds(self))

• Token IDs are passed as individual positional arguments to the C++ constructor (splatted with *token_ids).
• The token_ids property retrieves the list from the C++ side.

ImageData

ImageData wraps image tensor data along with an embedding size. It is registered as mlc.serve.ImageData.

@tvm_ffi.register_object("mlc.serve.ImageData")
class ImageData(Data):
    def __init__(self, image: Tensor, embed_size: int):
        self.embed_size = embed_size
        self.__init_handle_by_constructor__(_ffi_api.ImageData, image, embed_size)

Key methods:

• image property: Returns the underlying tvm.runtime.Tensor via FFI.
• __len__: Returns the embed size, providing a length interface compatible with prompt length calculations.
• from_url (static): Downloads or decodes an image from a URL (supports both HTTP URLs and base64-encoded data URIs), converts it to a NumPy array in NHWC format, and wraps it as a TVM tensor. Currently has hard-coded embed sizes: 576 by default and 1921 for phi3_v model type.
• get_embed_size (static): Computes the embedding size from model config as (image_size // patch_size) ** 2.
• get_input_size (static): Extracts the raw image size from the vision config.

SingleRequestStreamOutput

A Python @dataclass representing the streaming output for a single request.

@dataclass
class SingleRequestStreamOutput:
    delta_token_ids: List[int]
    delta_logprob_json_strs: Optional[List[str]]
    finish_reason: Optional[str]
    request_final_usage_json_str: Optional[str]
    extra_prefix_string: str

Fields:

• delta_token_ids: Newly generated token IDs since the last callback.
• delta_logprob_json_strs: Optional log probability JSON strings for the new tokens.
• finish_reason: The reason the request finished (e.g., "stop", "length"), or None if still generating.
• request_final_usage_json_str: Optional final usage statistics as a JSON string.
• extra_prefix_string: An additional prefix string attached to the output.

RequestStreamOutput

RequestStreamOutput is a TVM runtime object (registered as mlc.serve.RequestStreamOutput) that encapsulates the streamed delta output from the C++ engine. It is instantiated only on the C++ side.

@tvm_ffi.register_object("mlc.serve.RequestStreamOutput")
class RequestStreamOutput(Object):
    def unpack(self) -> Tuple[str, List[SingleRequestStreamOutput]]:
        fields = _ffi_api.RequestStreamOutputUnpack(self)
        request_final_usage_json_str = fields[4]
        request_id = str(fields[0])
        if request_final_usage_json_str is not None:
            return (
                request_id,
                [SingleRequestStreamOutput([], None, None, request_final_usage_json_str, "")],
            )
        # ... iterates over fields to build list of SingleRequestStreamOutput

unpack method logic:

1. Calls _ffi_api.RequestStreamOutputUnpack to get raw fields from the C++ side.
2. If request_final_usage_json_str (field index 4) is present, returns a single output with only usage data.
3. Otherwise, iterates over fields[1] (delta token IDs), fields[3] (finish reasons), and fields[5] (extra prefix strings) using zip, and optionally processes fields[2] (log probabilities) to construct a list of SingleRequestStreamOutput instances.

Class Hierarchy

Class	TVM FFI Name	Inherits From	Purpose
Data	mlc.serve.Data	tvm.runtime.Object	Base class for multi-modal data
TextData	mlc.serve.TextData	Data	Wraps a text string
TokenData	mlc.serve.TokenData	Data	Wraps a list of token IDs
ImageData	mlc.serve.ImageData	Data	Wraps an image tensor with embed size
RequestStreamOutput	mlc.serve.RequestStreamOutput	tvm.runtime.Object	Streamed output from C++ engine

Dependencies

• tvm, tvm_ffi: TVM runtime for object registration, tensor handling, and FFI bridging.
• PIL (Pillow): Used in ImageData.from_url for image decoding.
• requests: Used in ImageData.from_url for HTTP image fetching.
• numpy: Used in ImageData.from_url for array dimension expansion.

Design Notes

• All TVM-backed classes use __init_handle_by_constructor__ to delegate object construction to the C++ side, ensuring the Python objects are thin wrappers over C++ instances.
• The ImageData.from_url method contains a TODO comment noting that the embed size values (576 and 1921) are hard-coded for phi3.5-vision and llava models and should be computed dynamically.
• RequestStreamOutput has no Python-side constructor; it is created exclusively in C++ and unpacked on the Python side via the unpack() method.