feat: shard the SGLang worker gateway across processes (--gateway-workers)

components/src/dynamo/sglang/AGENTS.md

@@ -30,6 +30,35 @@ support the current version plus 1 version back (N and N-1). The pattern:
 component files. Do not version-check with `sglang.__version__` -- import probing is
 more reliable since SGLang's internal layout doesn't always match the version string.
 
+## Multi-process gateway (`--gateway-workers N`)
+
+`gateway.py`. One `dynamo.sglang` process fronts every DP rank of its engine: SGLang's
+`TokenizerManager` intake and the Dynamo handler's token relay run on one GIL. `--gateway-workers N` raises SGLang's `--tokenizer-worker-num` to N (that flag names the
+same process after only one of its jobs); with N > 1 SGLang puts a `MultiTokenizerRouter` in
+the engine process; it has no `generate_request`, so the leader (node rank 0) does not serve.
+Instead it publishes the launch data with SGLang's shared-memory contract
+(`write_data_for_multi_tokenizer`) and spawns N children (`--gateway-workers`, or `--tokenizer-worker-num` if only that is set),
+`python -m dynamo.sglang <same argv>` with `DYN_SGLANG_GATEWAY_PARENT_PID` set. A child
+builds a `TokenizerWorker` registered with the router (`build_gateway_engine`), wraps it
+in `GatewayEngine` (what the handlers use of an `sgl.Engine`: `tokenizer_manager`,
+`server_args`, `port_args`, `async_generate`, scheduler info) and runs the ordinary
+`init_decode`/`init_prefill` path as its own endpoint instance, so the router sees N
+instances per engine. `--gateway-workers` sets `tokenizer_worker_num` to exactly N so
+the router and child counts match; it is rejected for the direct-engine workers
+(embedding, rerank, multimodal, diffusion) and with `--enable-lora`, whose dynamic LoRA
+state would live in one child only.
+
+Metrics: the schedulers push KV metrics to one PULL socket and publish forward-pass
+metrics once, so child 0 (`DYN_SGLANG_GATEWAY_CHILD_INDEX=0`, `owns_engine_metrics()`)
+keeps the engine's `metrics_ipc_name` and the FPM relay; other children bind an idle
+private socket and advertise only the bootstrap zeros for their own instance. Every
+child republishes KV events under its own worker id, otherwise the router would see
+prefixes on instance 0 only and send it all traffic; the cost is N copies of each KV
+event on the event plane. The parent owns the engine subprocesses and the shared
+memory, runs the deferred shutdown handlers, and terminates and reaps the children.
+Snapshot mode and non-leader nodes are untouched. If SGLang exposes
+`Engine.attach_tokenizer_worker`, children use it instead of the facade.
+
 ## Entry Point
 
 `__main__.py` -> `main.py:main()` -> `main.py:worker()`

components/src/dynamo/sglang/backend_args.py

@@ -67,6 +67,19 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None:
             "the same SGLang-native pre/post processing with KV router support.",
         )
 
+        add_argument(
+            g,
+            flag_name="--gateway-workers",
+            env_var="DYN_SGL_GATEWAY_WORKERS",
+            default=None,
+            arg_type=int,
+            help="Run N gateway processes in front of this engine, each a Dynamo "
+            "endpoint instance with its own SGLang request gateway (SGLang calls "
+            "that process a tokenizer worker; it also handles request intake and "
+            "output relay, which is the work being spread). Implies SGLang's "
+            "--tokenizer-worker-num N unless that is set higher; setting only "
+            "--tokenizer-worker-num N runs N gateways as well.",
+        )
         add_negatable_bool_argument(
             g,
             flag_name="--enable-multimodal",
@@ -187,6 +200,7 @@ class DynamoSGLangConfig(ConfigBase):
     """Configuration for Dynamo SGLang wrapper (SGLang-specific only)."""
 
     use_sglang_tokenizer: bool
+    gateway_workers: Optional[int] = None
     # Internal roles derived from the canonical multimodal arguments in args.py.
     multimodal_encode_worker: bool = False
     multimodal_worker: bool = False
@@ -234,6 +248,8 @@ def validate(self) -> None:
                 "Both 'disagg_config' and 'disagg_config_key' must be provided together."
             )
 
+        if self.gateway_workers is not None and self.gateway_workers < 1:
+            raise ValueError("--gateway-workers must be a positive integer")
         self.validate_multimodal_topology()
 
         self.validate_dedicated_mm_encoder()

components/src/dynamo/sglang/gateway.py

@@ -0,0 +1,188 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-process request gateway for the SGLang worker: the leader keeps the engine
+and spawns N ``dynamo.sglang`` children that each serve requests through their own
+SGLang ``TokenizerWorker``. Design notes: AGENTS.md, "Multi-process gateway".
+"""
+
+from __future__ import annotations
+
+import asyncio
+import copy
+import logging
+import os
+import subprocess
+import sys
+import tempfile
+import types
+
+import sglang as sgl
+
+ENV_PARENT_PID = "DYN_SGLANG_GATEWAY_PARENT_PID"
+ENV_CHILD_INDEX = "DYN_SGLANG_GATEWAY_CHILD_INDEX"
+
+
+class GatewayEngine:
+    """What the request handlers need from an ``sgl.Engine``.
+
+    A child never owns scheduler processes; it only holds a ``TokenizerWorker`` bound
+    to the parent's router. ``Engine.async_generate`` only touches
+    ``self.tokenizer_manager``, so it can be reused unchanged.
+    """
+
+    def __init__(self, tokenizer_manager, server_args, port_args, scheduler_info):
+        self.tokenizer_manager = tokenizer_manager
+        self.server_args = server_args
+        self.port_args = port_args
+        self._scheduler_init_result = types.SimpleNamespace(
+            scheduler_infos=[scheduler_info]
+        )
+
+    async_generate = sgl.Engine.async_generate
+    _resolve_routed_dp_rank = sgl.Engine._resolve_routed_dp_rank
+
+    def shutdown(self):
+        pass
+
+
+def is_gateway_child() -> bool:
+    return ENV_PARENT_PID in os.environ
+
+
+def gateway_child_index() -> int:
+    return int(os.environ.get(ENV_CHILD_INDEX, "0"))
+
+
+def owns_engine_metrics() -> bool:
+    """SGLang's schedulers push KV metrics to one PULL socket and publish forward-pass
+    metrics once, so exactly one gateway process may consume them: child 0 (or the
+    single worker when the mode is off)."""
+    return gateway_child_index() == 0
+
+
+def gateway_worker_count(server_args, dynamo_args) -> int:
+    """Number of gateway processes to run for this engine; 1 disables the mode."""
+    if getattr(server_args, "tokenizer_worker_num", 1) <= 1:
+        return 1
+    if (getattr(server_args, "node_rank", 0) or 0) != 0:
+        return 1
+    if is_gateway_child():
+        return 1
+    n = getattr(dynamo_args, "gateway_workers", None)
+    return n if n else server_args.tokenizer_worker_num
+
+
+def _private_metrics_ipc(port_args):
+    # Only one process may bind the PULL socket the schedulers push KV metrics to;
+    # every other child gets an idle private one so its publisher can still start.
+    port_args = copy.copy(port_args)
+    port_args.metrics_ipc_name = (
+        f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}"
+    )
+    return port_args
+
+
+def _child_port_args(port_args):
+    return port_args if owns_engine_metrics() else _private_metrics_ipc(port_args)
+
+
+def build_gateway_engine():
+    """Child side: join the parent's engine through a ``TokenizerWorker``."""
+    parent_pid = int(os.environ[ENV_PARENT_PID])
+    attach = getattr(sgl.Engine, "attach_tokenizer_worker", None)
+    if attach is not None:
+        engine = attach(parent_pid)
+        engine.port_args = _child_port_args(engine.port_args)
+        logging.info(
+            "gateway child pid=%d attached via Engine.attach_tokenizer_worker",
+            os.getpid(),
+        )
+        return engine
+
+    from sglang.srt.managers.multi_tokenizer_mixin import (
+        get_tokenizer_worker_class,
+        read_from_shared_memory,
+    )
+    from sglang.srt.runtime_context import publish
+
+    port_args, server_args, scheduler_info = read_from_shared_memory(
+        f"multi_tokenizer_args_{parent_pid}"
+    )
+    publish(server_args, role="tokenizer")
+    port_args = _child_port_args(port_args)
+    port_args.tokenizer_ipc_name = (
+        f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}"
+    )
+    tm = get_tokenizer_worker_class(server_args)(server_args, port_args)
+    tm.max_req_input_len = scheduler_info["max_req_input_len"]
+    tm.set_startup_time(scheduler_info["startup_time"])
+    logging.info(
+        "gateway child pid=%d registered TokenizerWorker ipc=%s",
+        os.getpid(),
+        port_args.tokenizer_ipc_name,
+    )
+    return GatewayEngine(tm, server_args, port_args, scheduler_info)
+
+
+def _reap(proc: subprocess.Popen, timeout: float = 60) -> None:
+    try:
+        proc.wait(timeout=timeout)
+    except subprocess.TimeoutExpired:
+        proc.kill()
+        proc.wait()
+
+
+async def serve_via_gateway_children(
+    engine, count: int, shutdown_event: asyncio.Event
+) -> None:
+    """Parent side: publish the engine's launch data, spawn ``count`` children and
+    keep the engine alive until shutdown or until a child dies."""
+    from sglang.srt.managers.multi_tokenizer_mixin import write_data_for_multi_tokenizer
+
+    shm = getattr(engine, "_multi_tokenizer_shm", None)
+    owns_shm = shm is None
+    if shm is None:
+        scheduler_info = {
+            **engine._scheduler_init_result.scheduler_infos[0],
+            "startup_time": engine.tokenizer_manager.startup_time,
+        }
+        shm = write_data_for_multi_tokenizer(
+            engine.port_args, engine.server_args, scheduler_info
+        )
+    argv = sys.argv[1:]
+    procs: list[subprocess.Popen] = []
+    try:
+        for index in range(count):
+            env = {
+                **os.environ,
+                ENV_PARENT_PID: str(os.getpid()),
+                ENV_CHILD_INDEX: str(index),
+            }
+            procs.append(
+                subprocess.Popen(
+                    [sys.executable, "-m", "dynamo.sglang", *argv], env=env
+                )
+            )
+        logging.info(
+            "gateway parent pid=%d spawned %d children: %s",
+            os.getpid(),
+            count,
+            [p.pid for p in procs],
+        )
+        while not shutdown_event.is_set():
+            await asyncio.sleep(2)
+            dead = [p for p in procs if p.poll() is not None]
+            if dead:
+                raise RuntimeError(
+                    f"gateway child pid={dead[0].pid} exited rc={dead[0].returncode}"
+                )
+    finally:
+        for p in procs:
+            if p.poll() is None:
+                p.terminate()
+        await asyncio.gather(*(asyncio.to_thread(_reap, p) for p in procs))
+        if owns_shm:
+            try:
+                shm.unlink()
+            except FileNotFoundError:
+                pass

components/src/dynamo/sglang/init_llm.py

@@ -15,6 +15,7 @@
 from dynamo.llm import ModelInput, ModelType, WorkerType
 from dynamo.runtime import DistributedRuntime
 from dynamo.sglang.args import Config
+from dynamo.sglang.gateway import gateway_worker_count, serve_via_gateway_children
 from dynamo.sglang.health_check import (
     SglangDisaggHealthCheckPayload,
     SglangHealthCheckPayload,
@@ -77,6 +78,16 @@ async def init_decode(
         load_time = time.time() - start_time
 
     server_args = config.use_resolved_server_args(engine.server_args)
+    gateway_count = gateway_worker_count(server_args, dynamo_args)
+    if snapshot_engine is None and gateway_count > 1:
+        # engine.tokenizer_manager is SGLang's MultiTokenizerRouter here and cannot
+        # serve requests; gateway children do, this process keeps the engine alive.
+        try:
+            await serve_via_gateway_children(engine, gateway_count, shutdown_event)
+        finally:
+            if run_deferred_handlers is not None:
+                await run_deferred_handlers()
+        return
 
     if server_args.enable_trace:
         set_global_trace_level(dynamo_args.sglang_trace_level)
@@ -236,6 +247,16 @@ async def init_prefill(
         load_time = time.time() - start_time
 
     server_args = config.use_resolved_server_args(engine.server_args)
+    gateway_count = gateway_worker_count(server_args, dynamo_args)
+    if snapshot_engine is None and gateway_count > 1:
+        # engine.tokenizer_manager is SGLang's MultiTokenizerRouter here and cannot
+        # serve requests; gateway children do, this process keeps the engine alive.
+        try:
+            await serve_via_gateway_children(engine, gateway_count, shutdown_event)
+        finally:
+            if run_deferred_handlers is not None:
+                await run_deferred_handlers()
+        return
 
     if server_args.enable_trace:
         set_global_trace_level(dynamo_args.sglang_trace_level)

components/src/dynamo/sglang/main.py

@@ -19,6 +19,7 @@
 from dynamo.runtime.logging import configure_dynamo_logging
 from dynamo.sglang._compat import override_server_args
 from dynamo.sglang.args import parse_args
+from dynamo.sglang.gateway import build_gateway_engine, is_gateway_child
 from dynamo.sglang.init_diffusion import (
     init_image_diffusion,
     init_llm_diffusion,
@@ -62,12 +63,51 @@ async def worker(argv: list[str] | None = None):
             load_format=setup_gms(config.server_args),
         )
 
+    gateway_workers = getattr(config.dynamo_args, "gateway_workers", None) or 1
+    if gateway_workers > 1:
+        direct_engine_worker = [
+            flag
+            for flag in (
+                "image_diffusion_worker",
+                "video_generation_worker",
+                "rerank_worker",
+                "embedding_worker",
+                "multimodal_encode_worker",
+                "multimodal_worker",
+                "diffusion_worker",
+            )
+            if getattr(config.dynamo_args, flag, False)
+        ]
+        if direct_engine_worker:
+            raise ValueError(
+                "--gateway-workers > 1 is only supported by the decode and prefill "
+                f"LLM workers, not with --{direct_engine_worker[0].replace('_', '-')}"
+            )
+        if getattr(config.server_args, "enable_lora", False):
+            raise ValueError(
+                "--gateway-workers > 1 is not supported with --enable-lora: dynamic "
+                "LoRA state lives in each gateway process"
+            )
+        if config.server_args.tokenizer_worker_num != gateway_workers:
+            # SGLang launches its MultiTokenizerRouter (the shared-memory handoff the
+            # gateway children join) with exactly tokenizer_worker_num workers, so the
+            # router and child counts must match.
+            override_server_args(
+                config.server_args,
+                "dynamo.gateway",
+                tokenizer_worker_num=gateway_workers,
+            )
+
     # Snapshot mode: engine must be created before runtime so CRIU captures no
     # NATS/etcd connections.
     snapshot_controller = await prepare_snapshot_engine(config)
 
     dynamo_args = config.dynamo_args
     snapshot_engine = None
+    if snapshot_controller is None and is_gateway_child():
+        # Spawned by gateway.serve_via_gateway_children: share the parent's engine
+        # through a TokenizerWorker instead of launching schedulers.
+        snapshot_engine = build_gateway_engine()
     if snapshot_controller is not None:
         snapshot_engine = snapshot_controller.engine
         dynamo_args = await refresh_snapshot_restore_config(

components/src/dynamo/sglang/publisher.py

@@ -33,6 +33,7 @@
     local_dp_rank_bounds,
     publishes_kv_events,
 )
+from dynamo.sglang.gateway import gateway_child_index, owns_engine_metrics
 
 
 def get_local_dp_rank_range(server_args) -> range:
@@ -558,9 +559,18 @@ async def _idle() -> None:
 
     publisher.init_engine_metrics_publish()
     node_rank = getattr(config.server_args, "node_rank", 0) or 0
+    # Every gateway child republishes KV events under its own worker id so the
+    # router can match prefixes on any instance of the engine; scheduler metrics
+    # and forward-pass metrics are consumed once, by the process that owns them.
     if node_rank <= 0 and config.dynamo_args.use_kv_events:
         publisher.init_kv_event_publish()
-    publisher.init_fpm_relay()
+    if owns_engine_metrics():
+        publisher.init_fpm_relay()
+    else:
+        logging.info(
+            "gateway child %d: scheduler and forward-pass metrics are relayed by child 0",
+            gateway_child_index(),
+        )
 
     task = asyncio.create_task(publisher.run())
     logging.info("SGLang metrics loop started")