Add Mooncake store support for vLLM backend

Extends mooncake_master orchestration (previously SGLang-only) to the
vLLM backend so vLLM workers can use the in-process MooncakeStore
connector for cross-process KV sharing.

What's new
----------

* `VLLMMooncakeKVStoreConfig` schema (`backends/vllm.py`): mirrors the
  SGLang `mooncake_kv_store:` block, plus a vLLM-specific
  `store_config:` section. srtslurm renders that section into the
  `MOONCAKE_CONFIG_PATH` JSON file vLLM's `MooncakeStoreConnector`
  reads on startup (`MooncakeStoreConfig.load_from_env()`); the
  `env:` map is injected on every vLLM worker for the in-process
  `MC_*` knobs (e.g. `MC_ENABLE_DEST_DEVICE_AFFINITY`,
  `MC_STORE_CLIENT_METRIC`, `MC_TE_METRIC`).

* `start_mooncake_master` in `cli/do_sweep.py` now fires for vLLM as
  well as SGLang, writes the per-job store JSON next to the other
  log artifacts, and stamps the master RPC + HTTP-metadata endpoints
  on each worker. Shared launch constants moved to a new
  `backends/mooncake.py` (the SGLang module re-exports them so
  downstream imports keep working).

* Mooncake admin HTTP server (`/metrics`, `/health`, `/role`,
  `/ha_status`, `/leader`, `/query_key`) is now wired up explicitly:
  the master srun passes `--enable_metric_reporting=true
  --metrics_port=9003` (upstream default in master.cpp), and
  `start_mooncake_master` waits on the metrics port like it already
  waits on RPC + HTTP-metadata. The flag toggles a periodic stdout
  log thread only — `MasterAdminServer` listens unconditionally —
  the new constant's comment documents that to prevent future
  reordering.

* Docs (`docs/mooncake-kv-store.md`): adds a vLLM quick start, a
  vLLM-specific configuration reference (including `store_config`),
  expands the ownership table for the auto-stamped vars, and
  documents the new master metrics endpoint.

Validation
----------

* Manual: launched a 1P/1D disagg job; master came up with
  `enable_metric_reporting=1, metrics_port=9003`; vLLM workers
  picked up `MOONCAKE_CONFIG_PATH` and registered RDMA segments.
  `curl :9003/metrics` returned a full Prometheus dump.

* Tests: `make check` (ruff + pytest). New cases in
  `tests/test_e2e.py` cover the vLLM mooncake env injection and JSON
  rendering; `tests/test_dry_run.py` cases cover the new config
  surface in `srtctl dry-run` output.

Signed-off-by: inf-yasong <yasong.wang@inferact.ai>

Author: esmeetu

docs/mooncake-kv-store.md

@@ -1,13 +1,15 @@
 # Mooncake KV Store
 
-First-class support for [Mooncake](https://github.com/kvcache-ai/Mooncake) as the KV transfer backend for SGLang prefill-decode disaggregation. When `mooncake_kv_store` is set under an SGLang backend, srtslurm launches and configures the mooncake master automatically and wires up worker env vars so peer-to-peer transfers work across multiple nodes.
+First-class support for [Mooncake](https://github.com/kvcache-ai/Mooncake) as the KV transfer backend for prefill-decode disaggregation. When `mooncake_kv_store` is set under an SGLang or vLLM backend, srtslurm launches and configures the mooncake master automatically and wires up worker env vars so peer-to-peer transfers work across multiple nodes.
 
 ## Table of Contents
 
 - [Overview](#overview)
-- [Quick Start](#quick-start)
+- [Quick Start (SGLang)](#quick-start-sglang)
+- [Quick Start (vLLM)](#quick-start-vllm)
 - [What srtslurm Owns vs What You Set](#what-srtslurm-owns-vs-what-you-set)
 - [Configuration Reference](#configuration-reference)
+- [Master Metrics Endpoint](#master-metrics-endpoint)
 - [Validation](#validation)
 - [Common Configurations](#common-configurations)
   - [RDMA / InfiniBand](#rdma--infiniband)
@@ -30,7 +32,7 @@ Without first-class support, running mooncake with srtslurm meant:
 
 The `mooncake_kv_store` block automates 1–3. You still set the SGLang flags in step 4 because they're SGLang's CLI surface, not srtslurm's — but srtslurm validates that you did.
 
-## Quick Start
+## Quick Start (SGLang)
 
 Minimum config to run mooncake:
 
@@ -63,19 +65,53 @@ backend:
       disaggregation-transfer-backend: mooncake
 ```
 
+## Quick Start (vLLM)
+
+vLLM's `MooncakeStoreConnector` reads its configuration from a JSON file pointed to by `MOONCAKE_CONFIG_PATH` rather than directly from env vars, so the vLLM block takes an extra `store_config:` section that srtslurm renders into that JSON at job start:
+
+```yaml
+backend:
+  type: vllm
+  mooncake_kv_store:
+    env:                                  # injected on every vLLM worker
+      MOONCAKE_PROTOCOL: rdma
+      MC_ENABLE_DEST_DEVICE_AFFINITY: "1"
+    store_config:                         # → MOONCAKE_CONFIG_PATH JSON
+      metadata_server: "P2PHANDSHAKE"
+      global_segment_size: "50GB"
+      local_buffer_size: "4GB"
+      protocol: "rdma"
+      device_name: "mlx5_0,mlx5_1"
+  vllm_config:
+    prefill:
+      kv-transfer-config: '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}'
+    decode:
+      kv-transfer-config: '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}'
+```
+
+srtslurm stamps `MOONCAKE_MASTER`, `MOONCAKE_TE_META_DATA_SERVER`, `MOONCAKE_LOCAL_HOSTNAME`, and `MOONCAKE_CONFIG_PATH` on every worker; you supply the rest. `master_server_address` in `store_config` is also auto-filled from the infra node IP and ignored if set by hand.
+
+The `env:` map is injected on every vLLM worker (not on the standalone `mooncake_master` daemon — the master srun passes no env). Use it for in-process Mooncake C++ knobs like `MC_ENABLE_DEST_DEVICE_AFFINITY`, `MC_STORE_CLIENT_METRIC`, `MC_TE_METRIC`.
+
 ## What srtslurm Owns vs What You Set
 
 | Concern                                         | Owner     | Notes                                                                                                |
 | ----------------------------------------------- | --------- | ---------------------------------------------------------------------------------------------------- |
-| Launching `mooncake_master`                     | srtslurm  | Runs on the infra node (same node as etcd/nats; respects `infra.etcd_nats_dedicated_node`). Port 50051. |
+| Launching `mooncake_master`                     | srtslurm  | Runs on the infra node (same node as etcd/nats; respects `infra.etcd_nats_dedicated_node`). RPC `50051`, HTTP metadata `8080`, admin HTTP `9003`. |
 | `MOONCAKE_MASTER` env var on workers            | srtslurm  | Always computed as `<infra_node_ip>:50051`. User values in `env` are overridden.                      |
+| `MOONCAKE_TE_META_DATA_SERVER` env var          | srtslurm  | Always computed as `http://<infra_node_ip>:8080/metadata`.                                            |
 | `MOONCAKE_LOCAL_HOSTNAME` env var               | srtslurm  | Auto-resolved per-worker via `runtime.network_interface`. User can override in `env` for custom NICs. |
+| `MOONCAKE_CONFIG_PATH` (vLLM only)               | srtslurm  | Always points to the JSON file srtslurm renders from `store_config:`. Mounted under `/logs` in every worker. |
+| `master_server_address` in `store_config` (vLLM)| srtslurm  | Always overridden with `<infra_node_ip>:50051`. User values are ignored.                              |
 | `MOONCAKE_PROTOCOL`, `MOONCAKE_DEVICE`, etc.    | User      | Passed through `mooncake_kv_store.env` to all workers.                                               |
-| `disaggregation-transfer-backend: mooncake`     | User      | Set on `sglang_config.prefill` and `sglang_config.decode`. srtslurm validates this is present.       |
-| `disaggregation-ib-device`                      | User      | Set on `sglang_config.prefill` and `sglang_config.decode`. Format: `"mlx5_0,mlx5_1"` or JSON map.    |
+| `disaggregation-transfer-backend: mooncake`     | User      | (SGLang only) Set on `sglang_config.prefill` and `sglang_config.decode`. srtslurm validates this is present. |
+| `disaggregation-ib-device`                      | User      | (SGLang only) Set on `sglang_config.prefill` and `sglang_config.decode`. Format: `"mlx5_0,mlx5_1"` or JSON map. |
+| `kv-transfer-config`                            | User      | (vLLM only) Set on `vllm_config.prefill` and `vllm_config.decode` to wire vLLM's `MooncakeStoreConnector`. |
 
 ## Configuration Reference
 
+### SGLang
+
 ```yaml
 backend:
   type: sglang
@@ -93,10 +129,45 @@ backend:
       SGLANG_DISAGG_STAGING_POOL_SIZE_MB: "4096"
 ```
 
+### vLLM
+
+```yaml
+backend:
+  type: vllm
+  mooncake_kv_store:
+    container: ...                       # optional, default: job container
+    env:                                 # optional, injected on every vLLM worker
+      MOONCAKE_PROTOCOL: rdma
+      MC_ENABLE_DEST_DEVICE_AFFINITY: "1"
+      MC_STORE_CLIENT_METRIC: "1"        # default 1 (enabled)
+      MC_TE_METRIC: "0"                  # default 0 (disabled)
+    store_config:                        # optional, rendered into MOONCAKE_CONFIG_PATH JSON
+      metadata_server: "P2PHANDSHAKE"    # default "P2PHANDSHAKE"
+      global_segment_size: "4GB"         # default "4GB"
+      local_buffer_size: "4GB"           # default "4GB"
+      protocol: "rdma"                   # default "rdma"
+      device_name: "mlx5_0,mlx5_1"       # default ""
+```
+
 ### Fields
 
-- **`container`** (`str`, optional): Container image used for the `mooncake_master` srun. Defaults to the job container if unset. Useful when mooncake needs a different runtime than your SGLang container.
-- **`env`** (`dict[str, str]`, optional): Pass-through env vars injected on every prefill and decode worker. Keys map directly to mooncake's environment variable names — see the [SGLang server_args.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/environ.py) and [mooncake_store.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py) for the full list. Setting `MOONCAKE_MASTER` here is a no-op (srtslurm always wins).
+- **`container`** (`str`, optional): Container image used for the `mooncake_master` srun. Defaults to the job container if unset. Useful when mooncake needs a different runtime than your worker container.
+- **`env`** (`dict[str, str]`, optional): Pass-through env vars injected on every prefill and decode worker.
+  - For **SGLang**, keys map directly to mooncake's environment variable names — see the [SGLang server_args.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/environ.py) and [mooncake_store.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py) for the full list.
+  - For **vLLM**, this is for in-process Mooncake C++ knobs (`MC_*`) read by the transfer engine / store client. vLLM's connector itself reads configuration from `MOONCAKE_CONFIG_PATH` (the JSON rendered from `store_config:`), not from these env vars.
+  - Setting `MOONCAKE_MASTER`, `MOONCAKE_TE_META_DATA_SERVER`, or `MOONCAKE_CONFIG_PATH` here is a no-op (srtslurm always wins).
+- **`store_config`** (vLLM only, `dict[str, str]`, optional): Rendered as JSON into the file pointed to by `MOONCAKE_CONFIG_PATH`. Keys map 1:1 to vLLM's `MooncakeStoreConfig` dataclass. `master_server_address` is auto-filled and any user value is ignored.
+
+## Master Metrics Endpoint
+
+The `mooncake_master` admin HTTP server is always exposed on port `9003` on the infra node and starts before workers do (srtslurm waits for it). It serves:
+
+- `GET /metrics` — Prometheus text format (master + transfer-engine counters)
+- `GET /metrics/summary` — human-readable summary
+- `GET /health`, `/role`, `/ha_status`, `/leader`
+- `GET /query_key` — used by Dynamo's KV router shared-cache path
+
+To scrape from outside the cluster, point your collector at `http://<infra_node_ip>:9003/metrics`. The infra node IP is logged at job start.
 
 ## Validation
 

src/srtctl/backends/mooncake.py

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared mooncake_master constants used by both SGLang and vLLM backends.
+
+Kept in a dedicated module (rather than re-exported from one backend) so
+neither backend has to import from the other just to reach the port numbers.
+"""
+
+# RPC port the master listens on. Workers reach it via MOONCAKE_MASTER.
+MOONCAKE_MASTER_PORT = 50051
+
+# Port for the master's embedded HTTP metadata server (enabled with
+# --enable_http_metadata_server=true). Workers point MOONCAKE_TE_META_DATA_SERVER
+# at /metadata on this port so no separate metadata service is required.
+MOONCAKE_HTTP_METADATA_PORT = 8080
+
+# Port for the master's admin HTTP server. Matches the upstream default in
+# mooncake-store/src/master.cpp (--metrics_port=9003). Always listens once the
+# master is up — --enable_metric_reporting only toggles a periodic stdout log
+# thread, not the HTTP endpoints. Exposes /metrics (Prometheus text),
+# /metrics/summary, /health, /role, /ha_status, /leader, /query_key.
+MOONCAKE_METRICS_PORT = 9003

src/srtctl/backends/sglang.py

@@ -22,6 +22,10 @@
 from marshmallow import Schema
 from marshmallow_dataclass import dataclass
 
+# Re-exported so existing `from srtctl.backends.sglang import MOONCAKE_*` paths keep working.
+# Canonical home is srtctl.backends.mooncake.
+from srtctl.backends.mooncake import MOONCAKE_HTTP_METADATA_PORT, MOONCAKE_MASTER_PORT  # noqa: F401
+
 if TYPE_CHECKING:
     from srtctl.backends.base import SrunConfig
     from srtctl.core.runtime import RuntimeContext
@@ -30,9 +34,6 @@
 # Type alias for worker modes
 WorkerMode = Literal["prefill", "decode", "agg"]
 
-MOONCAKE_MASTER_PORT = 50051
-MOONCAKE_HTTP_METADATA_PORT = 8080
-
 
 @dataclass(frozen=True)
 class MooncakeKVStoreConfig:

src/srtctl/backends/vllm.py

@@ -24,14 +24,78 @@
 from marshmallow import Schema
 from marshmallow_dataclass import dataclass
 
+# vLLM reuses the same mooncake_master launch command and port pair as SGLang.
+# Disjoint port pairs can be reintroduced if we ever need to colocate an SGLang
+# and vLLM job on the same infra node.
+from srtctl.backends.mooncake import MOONCAKE_HTTP_METADATA_PORT, MOONCAKE_MASTER_PORT
+
 if TYPE_CHECKING:
     from srtctl.backends.base import SrunConfig
     from srtctl.core.runtime import RuntimeContext
-    from srtctl.core.topology import Endpoint, NodePortAllocator, Process
+    from srtctl.core.topology import Endpoint, Process
 
 # Type alias for worker modes
 WorkerMode = Literal["prefill", "decode", "agg"]
 
+# Filename for the mooncake-store JSON config srtslurm writes to log_dir at job
+# start. log_dir is mounted into every worker at /logs, so workers read the JSON
+# from MOONCAKE_STORE_CONFIG_CONTAINER_PATH.
+MOONCAKE_STORE_CONFIG_FILENAME = "mooncake_store_config.json"
+MOONCAKE_STORE_CONFIG_CONTAINER_PATH = f"/logs/{MOONCAKE_STORE_CONFIG_FILENAME}"
+
+
+@dataclass(frozen=True)
+class VLLMMooncakeKVStoreConfig:
+    """Mooncake KV store config for the vLLM backend.
+
+    When present, srtslurm launches ``mooncake_master`` on the infra node
+    (co-located with etcd/nats) using the shared SGLang launch command and
+    injects on every vLLM worker::
+
+        MOONCAKE_MASTER              = <infra_ip>:50051
+        MOONCAKE_TE_META_DATA_SERVER = http://<infra_ip>:8080/metadata
+        MOONCAKE_LOCAL_HOSTNAME      = <worker_ip>
+        MOONCAKE_CONFIG_PATH         = /logs/mooncake_store_config.json
+
+    The JSON file referenced by ``MOONCAKE_CONFIG_PATH`` is generated by
+    srtslurm at job start (see ``do_sweep.start_mooncake_master``) from
+    ``store_config`` below. vLLM's ``MooncakeStoreConnector`` reads this
+    file via ``MooncakeStoreConfig.load_from_env()``.
+
+    ``env:`` is injected on every vLLM worker (alongside the auto-stamped
+    ``MOONCAKE_*`` vars above), not on the standalone ``mooncake_master``
+    daemon — the master srun command passes no env. Use this for
+    in-process Mooncake C++ libraries linked into the worker:
+    ``MC_*`` knobs read by the transfer engine / store client
+    (e.g. ``MC_ENABLE_DEST_DEVICE_AFFINITY``, ``MC_STORE_CLIENT_METRIC``,
+    ``MC_TE_METRIC``), and any ``MOONCAKE_*`` overrides the connector
+    consults.
+
+    Example YAML::
+
+        backend:
+          type: vllm
+          mooncake_kv_store:
+            container: inferactinc/public:mk-int-20260507  # optional
+            env:                              # injected on every worker
+              MOONCAKE_PROTOCOL: rdma
+              MOONCAKE_GLOBAL_SEGMENT_SIZE: "4gb"
+              MOONCAKE_DEVICE: mlx5_0
+            store_config:                     # MooncakeStoreConfig JSON keys
+              metadata_server: "P2PHANDSHAKE"
+              global_segment_size: "100GB"
+              local_buffer_size: "4GB"
+              protocol: "rdma"
+              device_name: ""
+              # master_server_address: srtslurm auto-fills from infra IP
+    """
+
+    container: str | None = None
+    env: dict[str, str] = field(default_factory=dict)
+    store_config: dict[str, str] | None = None
+
+    Schema: ClassVar[builtins.type[Schema]] = Schema
+
 
 @dataclass(frozen=True)
 class VLLMServerConfig:
@@ -91,6 +155,11 @@ class VLLMProtocol:
     # dynamo 1.0.0+: translated to --kv-transfer-config (--connector was removed).
     connector: str | None = "nixl"
 
+    # Mooncake KV store — when set, srtslurm launches mooncake_master on the
+    # infra node and auto-injects MOONCAKE_MASTER / MOONCAKE_TE_META_DATA_SERVER
+    # / MOONCAKE_LOCAL_HOSTNAME on every vLLM worker.
+    mooncake_kv_store: VLLMMooncakeKVStoreConfig | None = None
+
     Schema: ClassVar[builtins.type[Schema]] = Schema
 
     # =========================================================================
@@ -132,14 +201,65 @@ def get_process_environment(self, process: Process) -> dict[str, str]:
         vLLM with dynamo requires unique ports for each worker:
         - DYN_VLLM_KV_EVENT_PORT: ZMQ port for KV events publishing
         - VLLM_NIXL_SIDE_CHANNEL_PORT: Port for NIXL side channel transfers
+        - VLLM_NIXL_SIDE_CHANNEL_HOST: Routable IP for NIXL side channel (not 0.0.0.0/localhost)
         """
+        from srtctl.core.slurm import get_hostname_ip
+
         env: dict[str, str] = {}
         if process.kv_events_port is not None:
             env["DYN_VLLM_KV_EVENT_PORT"] = str(process.kv_events_port)
         if process.nixl_port is not None:
             env["VLLM_NIXL_SIDE_CHANNEL_PORT"] = str(process.nixl_port)
+            env["VLLM_NIXL_SIDE_CHANNEL_HOST"] = get_hostname_ip(process.node)
         return env
 
+    def get_mooncake_worker_env(self, infra_node_ip: str, local_hostname: str) -> dict[str, str]:
+        """Get mooncake env vars to inject on a specific vLLM worker.
+
+        Returns ``{}`` when ``mooncake_kv_store`` is unset. Otherwise:
+
+        - ``MOONCAKE_MASTER`` and ``MOONCAKE_TE_META_DATA_SERVER`` are always
+          stamped by srtslurm (the user can't know the infra IP at config time).
+        - ``MOONCAKE_LOCAL_HOSTNAME`` defaults to the worker's resolved IP for
+          multi-node peer transfers, but a value in ``mooncake_kv_store.env``
+          wins (use this to pin to a specific RDMA NIC IP).
+        - ``MOONCAKE_CONFIG_PATH`` points to the JSON file srtslurm writes at
+          job start (mounted into the container at ``/logs``). vLLM's
+          ``MooncakeStoreConnector`` requires this — it does not read the
+          ``MOONCAKE_*`` env vars directly.
+        """
+        if self.mooncake_kv_store is None:
+            return {}
+        return {
+            "MOONCAKE_LOCAL_HOSTNAME": local_hostname,
+            **self.mooncake_kv_store.env,
+            "MOONCAKE_MASTER": f"{infra_node_ip}:{MOONCAKE_MASTER_PORT}",
+            "MOONCAKE_TE_META_DATA_SERVER": (f"http://{infra_node_ip}:{MOONCAKE_HTTP_METADATA_PORT}/metadata"),
+            "MOONCAKE_CONFIG_PATH": MOONCAKE_STORE_CONFIG_CONTAINER_PATH,
+        }
+
+    def build_mooncake_store_config(self, infra_node_ip: str) -> dict[str, str]:
+        """Build the JSON payload for vLLM's ``MooncakeStoreConfig.load_from_env()``.
+
+        Keys map 1:1 to vLLM's ``MooncakeStoreConfig`` dataclass. Values come
+        from ``mooncake_kv_store.store_config`` when set; missing keys fall back
+        to defaults. ``master_server_address`` is always auto-filled from the
+        infra node IP (any user-provided value is overridden — the user can't
+        know the infra IP at config time).
+        """
+        user_cfg: dict[str, str] = {}
+        if self.mooncake_kv_store is not None and self.mooncake_kv_store.store_config:
+            user_cfg = dict(self.mooncake_kv_store.store_config)
+
+        return {
+            "metadata_server": user_cfg.get("metadata_server", "P2PHANDSHAKE"),
+            "master_server_address": f"{infra_node_ip}:{MOONCAKE_MASTER_PORT}",
+            "global_segment_size": user_cfg.get("global_segment_size", "4GB"),
+            "local_buffer_size": user_cfg.get("local_buffer_size", "4GB"),
+            "protocol": user_cfg.get("protocol", "rdma"),
+            "device_name": user_cfg.get("device_name", ""),
+        }
+
     def get_served_model_name(self, default: str) -> str:
         """Get served model name from vLLM config, or return default."""
         if self.vllm_config:
@@ -193,7 +313,6 @@ def endpoints_to_processes(
         self,
         endpoints: list[Endpoint],
         base_sys_port: int = 8081,
-        port_allocator: NodePortAllocator | None = None,
     ) -> list[Process]:
         """Convert endpoints to processes.
 
@@ -207,13 +326,12 @@ def endpoints_to_processes(
 
         if not has_dp_mode:
             # Standard TP mode: one process per node
-            return endpoints_to_processes(endpoints, base_sys_port=base_sys_port, port_allocator=port_allocator)
+            return endpoints_to_processes(endpoints, base_sys_port=base_sys_port)
 
         # DP+EP mode: one process per GPU
         processes: list[Process] = []
         current_sys_port = base_sys_port
-        if port_allocator is None:
-            port_allocator = NodePortAllocator()
+        port_allocator = NodePortAllocator()
 
         for endpoint in endpoints:
             if not self._is_dp_mode(endpoint.mode):