From 8f0f36c64b09d439d7daf2676259419b46721c65 Mon Sep 17 00:00:00 2001
From: ishandhanani <82981111+ishandhanani@users.noreply.github.com>
Date: Wed, 18 Mar 2026 15:38:32 -0500
Subject: [PATCH] [1/2] Add ModelExpress coordination for remote instance
 weight loading - matching TP (#19920)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Ishan Dhanani <ishan@dhanani.dev>
---
 docs/advanced_features/rfork.md               |  33 ++++-
 python/sglang/srt/configs/load_config.py      |   2 +
 .../sglang/srt/model_executor/model_runner.py |  90 ++++++++++++
 python/sglang/srt/model_loader/loader.py      | 138 ++++++++++++++++++
 .../remote_instance_weight_loader_utils.py    |   1 +
 python/sglang/srt/server_args.py              |  67 ++++++++-
 6 files changed, 320 insertions(+), 11 deletions(-)

diff --git a/docs/advanced_features/rfork.md b/docs/advanced_features/rfork.md
index 5e01aa111..e4b513328 100644
--- a/docs/advanced_features/rfork.md
+++ b/docs/advanced_features/rfork.md
@@ -9,11 +9,12 @@ To learn more details about R-Fork, please check **<a href=https://lmsys.org/blo
 | Argument     | Usage                                      |
 |--------------|--------------------------------------------|
 | load-format  | set to `remote_instance` to enable R-Fork. |
-| remote-instance-weight-loader-backend | `nccl` or `transfer_engine`, default value is `nccl` |
-| remote-instance-weight-loader-seed-instance-ip | IP address of the seed instance who will provide the model weight |
-| remote-instance-weight-loader-seed-instance-service-port | the port that the seed instance's HTTP server is listening on |
-| remote-instance-weight-loader-send-weights-group-ports | the list of available ports on the seed instance that will be used to build NCCL communication groups between seed and client instance. This argument is only needed by `nccl` backend.  |
-| remote-instance-weight-loader-start-seed-via-transfer-engine | set to start seed service that supports TransferEngine as backend. It is needed for seed instances when using `transfer_engine` as backend. |
+| remote-instance-weight-loader-backend | `nccl`, `transfer_engine`, or `modelexpress`. Default is `nccl`. |
+| remote-instance-weight-loader-seed-instance-ip | IP address of the seed instance who will provide the model weight. Used by `nccl` and `transfer_engine` backends. |
+| remote-instance-weight-loader-seed-instance-service-port | the port that the seed instance's HTTP server is listening on. Used by `nccl` and `transfer_engine` backends. |
+| remote-instance-weight-loader-send-weights-group-ports | the list of available ports on the seed instance that will be used to build NCCL communication groups between seed and client instance. Only needed by `nccl` backend. |
+| remote-instance-weight-loader-start-seed-via-transfer-engine | set to start seed service that supports TransferEngine as backend. Needed for seed instances when using `transfer_engine` as backend. |
+| modelexpress-config | JSON config for `modelexpress` backend. Keys: `"url"` (required, gRPC host:port of ModelExpress server), `"model_name"` (optional, defaults to `--model-path`), `"source"` (optional bool, `true` for seed mode). |
 
 ### NCCL as backend
 
@@ -47,3 +48,25 @@ python -m sglang.launch_server [args] \
   --remote-instance-weight-loader-seed-instance-service-port [seed_instance_service_port] \
   --remote-instance-weight-loader-backend transfer_engine
 ```
+
+### ModelExpress as backend
+
+[ModelExpress](https://github.com/ai-dynamo/modelexpress) is a coordination service that manages P2P weight transfer metadata. It removes the need for direct seed IP/port configuration by providing a centralized registry that seeds publish to and clients discover from. Under the hood it uses TransferEngine (Mooncake) for the actual RDMA data transfer.
+
+A running ModelExpress server is required. See the [ModelExpress documentation](https://github.com/ai-dynamo/modelexpress) for setup instructions.
+
+seed instance:
+```shell
+python -m sglang.launch_server [args] \
+  --modelexpress-config '{"url": "[modelexpress_grpc_host:port]", "model_name": "[model_name]", "source": true}'
+```
+
+client instance:
+```shell
+python -m sglang.launch_server [args] \
+  --load-format remote_instance \
+  --remote-instance-weight-loader-backend modelexpress \
+  --modelexpress-config '{"url": "[modelexpress_grpc_host:port]", "model_name": "[model_name]"}'
+```
+
+The seed publishes its TransferEngine session ID and tensor layout to ModelExpress. The client queries ModelExpress to discover the seed, then pulls weights directly via RDMA. This enables dynamic seed discovery without hardcoding IPs, and supports multiple models through a single ModelExpress instance.
diff --git a/python/sglang/srt/configs/load_config.py b/python/sglang/srt/configs/load_config.py
index ddf8d2967..443ba643d 100644
--- a/python/sglang/srt/configs/load_config.py
+++ b/python/sglang/srt/configs/load_config.py
@@ -76,6 +76,8 @@ class LoadConfig:
     remote_instance_weight_loader_send_weights_group_ports: Optional[List[int]] = None
     remote_instance_weight_loader_backend: Optional[str] = None
     remote_instance_weight_loader_transfer_engine: Optional[Any] = None
+    modelexpress_url: Optional[str] = None
+    modelexpress_model_name: Optional[str] = None
 
     # ModelOpt-specific loading options
     modelopt_checkpoint_restore_path: Optional[str] = None
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index 02b856c6f..62079dacf 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -676,6 +676,74 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             local_ip, self.remote_instance_transfer_engine.get_rpc_port()
         ).to_host_port_str()
 
+    def _publish_modelexpress_metadata(self):
+        """Publish TransferEngine metadata to ModelExpress server (seed mode)."""
+        try:
+            from modelexpress import p2p_pb2
+            from modelexpress.client import MxClient
+        except ImportError as exc:
+            raise ImportError(
+                "ModelExpress support requires the 'modelexpress' package. "
+                "Install it with: pip install modelexpress"
+            ) from exc
+
+        model_name = (
+            self.server_args.modelexpress_model_name or self.server_args.model_path
+        )
+        mx_url = self.server_args.modelexpress_url
+        session_id = self.remote_instance_transfer_engine_session_id
+        weight_info = self.remote_instance_transfer_engine_weight_info
+
+        if not session_id or weight_info is None:
+            logger.warning(
+                "ModelExpress source: skipping publish -- "
+                "TransferEngine not initialized or no weight info"
+            )
+            return
+
+        # Build tensor descriptors from weight_info dict
+        tensors = []
+        for name, (addr, numel, element_size) in weight_info.items():
+            tensors.append(
+                p2p_pb2.TensorDescriptor(
+                    name=name,
+                    addr=addr,
+                    size=numel * element_size,
+                    device_id=self.gpu_id,
+                )
+            )
+
+        worker = p2p_pb2.WorkerMetadata(
+            worker_rank=self.tp_rank,
+            transfer_engine_session_id=session_id,
+            tensors=tensors,
+        )
+
+        mx_client = MxClient(server_url=mx_url)
+        try:
+            logger.info(
+                "ModelExpress source: publishing metadata for model=%s, "
+                "tp_rank=%d, session=%s, %d tensors",
+                model_name,
+                self.tp_rank,
+                session_id,
+                len(tensors),
+            )
+            mx_client.publish_metadata(model_name, [worker])
+            mx_client.publish_ready(
+                model_name,
+                worker_id=self.tp_rank,
+                session_id=mx_client.session_id,
+                metadata_hash="",
+            )
+            logger.info(
+                "ModelExpress source: published ready for model=%s, tp_rank=%d",
+                model_name,
+                self.tp_rank,
+            )
+        finally:
+            mx_client.close()
+
     def model_specific_adjustment(self):
         server_args = self.server_args
 
@@ -963,6 +1031,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             remote_instance_weight_loader_send_weights_group_ports=self.server_args.remote_instance_weight_loader_send_weights_group_ports,
             remote_instance_weight_loader_backend=self.server_args.remote_instance_weight_loader_backend,
             remote_instance_weight_loader_transfer_engine=self.remote_instance_transfer_engine,
+            modelexpress_url=self.server_args.modelexpress_url,
+            modelexpress_model_name=self.server_args.modelexpress_model_name
+            or self.server_args.model_path,
             modelopt_config=modelopt_config,
             rl_quant_profile=self.server_args.rl_quant_profile,
             draft_model_idx=self.draft_model_idx,
@@ -1015,6 +1086,25 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 )
         monkey_patch_vllm_parallel_state(reverse=True)
 
+        # Publish metadata to ModelExpress if running as seed source
+        if self.server_args.modelexpress_source:
+            # Seed loads via DefaultModelLoader (load_format=auto), which doesn't
+            # call register_memory_region(). Do it here so weight_info is populated.
+            if (
+                self.remote_instance_transfer_engine_weight_info is None
+                and self.remote_instance_transfer_engine is not None
+            ):
+                from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
+                    register_memory_region,
+                )
+
+                self.remote_instance_transfer_engine_weight_info = (
+                    register_memory_region(
+                        self.model, self.remote_instance_transfer_engine
+                    )
+                )
+            self._publish_modelexpress_metadata()
+
         get_offloader().post_init()
 
         # Register model for layerwise NVTX profiling if enabled
diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py
index 576d02e3a..3087f0e85 100644
--- a/python/sglang/srt/model_loader/loader.py
+++ b/python/sglang/srt/model_loader/loader.py
@@ -2151,6 +2151,15 @@ class RemoteInstanceModelLoader(BaseModelLoader):
                 raise RuntimeError(
                     "Failed to load weights from remote instance via transfer engine."
                 )
+        elif (
+            load_config.remote_instance_weight_loader_backend
+            == RemoteInstanceWeightLoaderBackend.MODELEXPRESS
+        ):
+            self.load_model_from_modelexpress(
+                model,
+                load_config,
+                device_config,
+            )
         else:
             raise ValueError("Invalid remote instance weight loader backend.")
 
@@ -2264,6 +2273,135 @@ class RemoteInstanceModelLoader(BaseModelLoader):
 
         return True
 
+    def load_model_from_modelexpress(
+        self,
+        model,
+        load_config: LoadConfig,
+        device_config: DeviceConfig,
+    ):
+        """Load weights via ModelExpress coordination + TransferEngine RDMA."""
+        try:
+            from modelexpress.client import MxClient
+        except ImportError as exc:
+            raise ImportError(
+                "ModelExpress support requires the 'modelexpress' package. "
+                "Install it with: pip install modelexpress"
+            ) from exc
+
+        transfer_engine = load_config.remote_instance_weight_loader_transfer_engine
+        if transfer_engine is None:
+            raise RuntimeError(
+                "TransferEngine is not initialized for modelexpress backend."
+            )
+        tp_rank = load_config.tp_rank
+        model_name = load_config.modelexpress_model_name
+
+        logger.info(
+            "ModelExpress: registering memory regions for tp_rank=%d...", tp_rank
+        )
+        self.remote_instance_transfer_engine_weight_info = register_memory_region(
+            model, transfer_engine
+        )
+
+        # Wait for seed to be ready via ModelExpress
+        mx_client = MxClient(server_url=load_config.modelexpress_url)
+        try:
+            logger.info(
+                "ModelExpress: waiting for seed ready (model=%s)...",
+                model_name,
+            )
+            ready, session_id, metadata_hash = mx_client.wait_for_ready(
+                model_name,
+                worker_id=tp_rank,
+            )
+            if not ready:
+                raise RuntimeError(
+                    f"ModelExpress: timed out waiting for seed ready "
+                    f"(model={model_name}, worker={tp_rank})"
+                )
+
+            response = mx_client.get_metadata(model_name)
+            if not response.found:
+                raise RuntimeError(
+                    f"ModelExpress: no metadata found for model={model_name}"
+                )
+
+            # Find the worker matching our tp_rank
+            source_worker = None
+            for w in response.workers:
+                if w.worker_rank == tp_rank:
+                    source_worker = w
+                    break
+            if source_worker is None:
+                raise RuntimeError(
+                    f"ModelExpress: no worker metadata for rank={tp_rank}"
+                )
+
+            # Extract session_id from oneof backend_metadata
+            backend_field = source_worker.WhichOneof("backend_metadata")
+            if backend_field == "transfer_engine_session_id":
+                seed_session_id = source_worker.transfer_engine_session_id
+            else:
+                raise RuntimeError(
+                    f"ModelExpress: expected transfer_engine_session_id, "
+                    f"got backend_metadata={backend_field}"
+                )
+
+            # Build {name: (addr, size_bytes)} from seed tensor descriptors
+            seed_weight_info = {}
+            for td in source_worker.tensors:
+                seed_weight_info[td.name] = (td.addr, td.size)
+
+            logger.info(
+                "ModelExpress: got %d tensor descriptors from seed (session=%s)",
+                len(seed_weight_info),
+                seed_session_id,
+            )
+        finally:
+            mx_client.close()
+
+        # Transfer weights via TransferEngine RDMA
+        seed_ptr_list = []
+        client_ptr_list = []
+        client_len_list = []
+        for name, tensor in model.named_parameters():
+            weight_info = seed_weight_info.get(name, None)
+            if weight_info is None:
+                raise RuntimeError(
+                    f"ModelExpress: cannot find weight info for {name} "
+                    f"in seed metadata"
+                )
+            seed_ptr, seed_size = weight_info
+            local_size = tensor.numel() * tensor.element_size()
+            if seed_size != local_size:
+                raise RuntimeError(
+                    f"ModelExpress: size mismatch for {name}: "
+                    f"seed={seed_size} bytes, local={local_size} bytes"
+                )
+            seed_ptr_list.append(seed_ptr)
+            client_ptr_list.append(tensor.data_ptr())
+            client_len_list.append(local_size)
+
+        logger.info(
+            "ModelExpress: starting RDMA transfer of %d tensors...",
+            len(seed_ptr_list),
+        )
+        ret = transfer_engine.batch_transfer_sync_read(
+            seed_session_id,
+            client_ptr_list,
+            seed_ptr_list,
+            client_len_list,
+        )
+        if ret < 0:
+            raise RuntimeError(
+                f"ModelExpress: batch_transfer_sync_read failed, error={ret}"
+            )
+
+        if hasattr(model, "post_load_weights"):
+            model.post_load_weights()
+
+        logger.info("ModelExpress: weight transfer complete for tp_rank=%d", tp_rank)
+
 
 class RemoteModelLoader(BaseModelLoader):
     """Model loader that can load Tensors from remote database."""
diff --git a/python/sglang/srt/model_loader/remote_instance_weight_loader_utils.py b/python/sglang/srt/model_loader/remote_instance_weight_loader_utils.py
index c063ea342..8a945bb4c 100644
--- a/python/sglang/srt/model_loader/remote_instance_weight_loader_utils.py
+++ b/python/sglang/srt/model_loader/remote_instance_weight_loader_utils.py
@@ -15,6 +15,7 @@ logger = logging.getLogger(__name__)
 class RemoteInstanceWeightLoaderBackend(str, enum.Enum):
     NCCL = "nccl"
     TRANSFER_ENGINE = "transfer_engine"
+    MODELEXPRESS = "modelexpress"
 
 
 def trigger_init_weights_send_group_for_remote_instance_request(
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index 448c7d6c4..bad3d2529 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -700,8 +700,11 @@ class ServerArgs:
     remote_instance_weight_loader_seed_instance_ip: Optional[str] = None
     remote_instance_weight_loader_seed_instance_service_port: Optional[int] = None
     remote_instance_weight_loader_send_weights_group_ports: Optional[List[int]] = None
-    remote_instance_weight_loader_backend: Literal["transfer_engine", "nccl"] = "nccl"
+    remote_instance_weight_loader_backend: Literal[
+        "transfer_engine", "nccl", "modelexpress"
+    ] = "nccl"
     remote_instance_weight_loader_start_seed_via_transfer_engine: bool = False
+    modelexpress_config: Optional[str] = None
 
     # For PD-Multiplexing
     enable_pdmux: bool = False
@@ -2967,7 +2970,19 @@ class ServerArgs:
             self.custom_weight_loader = []
 
         if self.load_format == "remote_instance":
-            if (
+            if self.remote_instance_weight_loader_backend == "modelexpress":
+                # ModelExpress backend: requires url in --modelexpress-config
+                if self.modelexpress_url is None:
+                    logger.warning(
+                        "Fallback load_format to 'auto' due to missing 'url' in --modelexpress-config."
+                    )
+                    self.load_format = "auto"
+                elif not self.validate_transfer_engine():
+                    logger.warning(
+                        "Fallback load_format to 'auto' due to 'transfer_engine' (required by modelexpress) not being supported."
+                    )
+                    self.load_format = "auto"
+            elif (
                 self.remote_instance_weight_loader_seed_instance_ip is None
                 or self.remote_instance_weight_loader_seed_instance_service_port is None
             ):
@@ -5567,15 +5582,21 @@ class ServerArgs:
         parser.add_argument(
             "--remote-instance-weight-loader-backend",
             type=str,
-            choices=["transfer_engine", "nccl"],
+            choices=["transfer_engine", "nccl", "modelexpress"],
             default=ServerArgs.remote_instance_weight_loader_backend,
-            help="The backend for loading weights from remote instance. Can be 'transfer_engine' or 'nccl'. Default is 'nccl'.",
+            help="The backend for loading weights from remote instance. Can be 'transfer_engine', 'nccl', or 'modelexpress'. Default is 'nccl'.",
         )
         parser.add_argument(
             "--remote-instance-weight-loader-start-seed-via-transfer-engine",
             action="store_true",
             help="Start seed server via transfer engine backend for remote instance weight loader.",
         )
+        parser.add_argument(
+            "--modelexpress-config",
+            type=str,
+            default=ServerArgs.modelexpress_config,
+            help='JSON config for ModelExpress P2P weight loading. Keys: "url" (required, gRPC host:port), "model_name" (optional, defaults to --model-path), "source" (optional bool, true for seed mode). Example: \'{"url": "localhost:8001", "model_name": "my-model", "source": true}\'',
+        )
 
         # For PD-Multiplexing
         parser.add_argument(
@@ -6109,7 +6130,11 @@ class ServerArgs:
         )
 
     def validate_transfer_engine(self):
-        if importlib.util.find_spec("mooncake.engine") is None:
+        try:
+            mooncake_available = importlib.util.find_spec("mooncake.engine") is not None
+        except (ModuleNotFoundError, ValueError):
+            mooncake_available = False
+        if not mooncake_available:
             logger.warning(
                 "Failed to import mooncake.engine. Does not support using TransferEngine as remote instance weight loader backend."
             )
@@ -6122,14 +6147,44 @@ class ServerArgs:
         else:
             return True
 
+    @property
+    def _parsed_modelexpress_config(self) -> dict:
+        cache = getattr(self, "_mx_config_cache", None)
+        if cache is not None:
+            return cache
+        if self.modelexpress_config is None:
+            result = {}
+        elif isinstance(self.modelexpress_config, str):
+            result = json.loads(self.modelexpress_config)
+        else:
+            result = self.modelexpress_config
+        object.__setattr__(self, "_mx_config_cache", result)
+        return result
+
+    @property
+    def modelexpress_url(self) -> Optional[str]:
+        return self._parsed_modelexpress_config.get("url")
+
+    @property
+    def modelexpress_model_name(self) -> Optional[str]:
+        return self._parsed_modelexpress_config.get("model_name")
+
+    @property
+    def modelexpress_source(self) -> bool:
+        return self._parsed_modelexpress_config.get("source", False)
+
     def remote_instance_weight_loader_use_transfer_engine(self):
         # Use TransferEngine as seed backend.
         if self.remote_instance_weight_loader_start_seed_via_transfer_engine:
             return True
+        # ModelExpress source mode also needs TransferEngine init.
+        if self.modelexpress_source:
+            return True
         # Use TransferEngine as client backend.
         elif (
             self.load_format == "remote_instance"
-            and self.remote_instance_weight_loader_backend == "transfer_engine"
+            and self.remote_instance_weight_loader_backend
+            in ("transfer_engine", "modelexpress")
         ):
             return True
         else: