Implementation:Mlc ai Mlc llm Router Init

Knowledge Sources	MLC-LLM
Domains	Deep_Learning, Distributed_Serving
Last Updated	2026-02-09 00:00 GMT

Overview

Concrete tool for setting up distributed inference infrastructure with multiple engine instances across GPUs, provided by MLC-LLM.

Description

Router.__init__ spawns multiple MLC-LLM engine server processes (one per endpoint), each bound to a distinct set of GPU devices. It initializes NVSHMEM for GPU-direct KV cache transfers across engines, computes device ID partitions for each endpoint, and concurrently starts all servers using background threads. On completion, the router holds references to all server processes, their URLs, load counters, and a tokenizer for prompt encoding.

The constructor performs the following sequence:

1. Validates that the lengths of hosts, ports, and num_gpus are equal.
2. Calls tvm.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") to generate a shared NVSHMEM UID.
3. Computes cumulative GPU device ID offsets (device_id_starts) so that engine i uses GPUs starting at device_id_starts[i].
4. Spawns each PopenServer in a separate thread, passing the NVSHMEM configuration (UID, total PEs, PE start) via the MLC_NVSHMEM_INIT_CONFIG_JSON_STR environment variable.
5. Waits for all threads to complete, then initializes the tokenizer.

Usage

Use Router.__init__ when you need to programmatically create a multi-engine disaggregated serving deployment. This is the entry point for setting up the full inference infrastructure before handling any requests. It is typically called indirectly through the serve() function in mlc_llm.interface.router, but can also be used directly for custom router implementations.

Code Reference

Source Location

• Repository: MLC-LLM
• File: python/mlc_llm/router/router.py (Lines 20-104)

Signature

class Router:
    def __init__(
        self,
        model: str,
        model_lib: Optional[str] = None,
        hosts: Optional[List[str]] = None,
        ports: Optional[List[int]] = None,
        num_gpus: Optional[List[int]] = None,
        enable_prefix_cache: bool = False,
        router_mode: Literal["disagg", "round-robin"] = "disagg",
        pd_balance_factor: float = 0.0,
    ):

Import

from mlc_llm.router import Router

I/O Contract

Inputs

Name	Type	Required	Description
model	str	Yes	Path or identifier of the model to serve. Passed to each PopenServer and used to initialize the tokenizer.
model_lib	Optional[str]	No	Path to the compiled model library. Defaults to None (auto-detected).
hosts	Optional[List[str]]	No	List of hostnames/IPs for each endpoint. Defaults to ["127.0.0.1"].
ports	Optional[List[int]]	No	List of port numbers for each endpoint. Defaults to [8080].
num_gpus	Optional[List[int]]	No	Number of GPUs for each endpoint. Defaults to [1]. Must have same length as hosts and ports.
enable_prefix_cache	bool	No	Whether to enable radix-tree prefix caching. Defaults to False. When enabled, sets prefix_cache_mode="radix" in the engine config.
router_mode	Literal["disagg", "round-robin"]	No	Routing strategy. "disagg" separates prefill and decode across engines; "round-robin" dispatches full requests to least-loaded engines. Defaults to "disagg".
pd_balance_factor	float	No	Controls the split point between prefill and decode in disaggregated mode. A value of 0.0 means the entire prompt is prefilled on the prefill engine. Higher values shift more work to the decode engine. Defaults to 0.0.

Outputs

Name	Type	Description
(instance)	Router	A fully initialized Router with all engine endpoints running, NVSHMEM configured, and tokenizer loaded. Key instance attributes: servers (list of PopenServer), server_urls (list of endpoint URLs), num_running_requests (per-endpoint load counters), tokenizer (the loaded tokenizer).

Usage Examples

Basic Usage

from mlc_llm.router import Router

# Create a disaggregated router with 2 endpoints:
# - Engine 0 (prefill): GPU 0, port 8080
# - Engine 1 (decode): GPU 1, port 8081
router = Router(
    model="dist/Llama-2-7b-chat-hf-q4f16_1-MLC",
    hosts=["127.0.0.1", "127.0.0.1"],
    ports=[8080, 8081],
    num_gpus=[1, 1],
    router_mode="disagg",
)

# Use the router to handle requests...
# ...

# Clean up when done
router.terminate()

Multi-GPU Endpoints with Prefix Cache

from mlc_llm.router import Router

# Create a router with 3 endpoints using tensor parallelism:
# - Engine 0 (prefill): GPUs 0-1, port 8080
# - Engine 1 (decode): GPUs 2-3, port 8081
# - Engine 2 (decode): GPUs 4-5, port 8082
router = Router(
    model="dist/Llama-2-70b-chat-hf-q4f16_1-MLC",
    hosts=["127.0.0.1", "127.0.0.1", "127.0.0.1"],
    ports=[8080, 8081, 8082],
    num_gpus=[2, 2, 2],
    enable_prefix_cache=True,
    router_mode="disagg",
    pd_balance_factor=0.0,
)

Related Pages

Implements Principle

• Principle:Mlc_ai_Mlc_llm_Multi_Engine_Infrastructure