From 228e36998fb0a6514d32f7af8568a109a3f98780 Mon Sep 17 00:00:00 2001 From: xyuzh Date: Tue, 9 Jun 2026 16:16:35 -0700 Subject: [PATCH 01/10] [sglang-miles] RDT/NIXL weight sync support for Ray scheduler actors Add the sglang-side support for RDT (Ray Direct Transport / NIXL) weight sync, used by the miles trainer to push weights to rollout engines via a zero-copy RDMA pull instead of NCCL broadcast. - scheduler_actor: add pull_weights(), which uses ray.experimental set_target_for_ref to RDMA pre-sharded weight buckets directly into the model's param.data buffers (no intermediate receive buffers / copies). - ray/engine: register SchedulerActors as detached named actors with the http port baked into the name, so the trainer (a different Ray job) can discover them via list_named_actors even when several engines are co-located on one node; raise max_concurrency so a concurrent pull_weights is not starved while run_event_loop blocks; set RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1 so the absolute GPU id from get_accelerator_ids() stays valid. - server_args: add enable_engine_info_bootstrap, plus the needs_engine_info_bootstrap() / registers_parallelism_config() predicates that fold it in alongside the existing transfer-engine conditions. - engine / model_runner: gate the bootstrap server and the parallelism-config registration on those predicates, so RDT can get /parallelism_config WITHOUT the mooncake/verbs P2P transfer-engine seeding. Requires ray>=2.55.1 for ray.experimental.set_target_for_ref, which is already the floor in python/pyproject.toml. Co-Authored-By: Claude Opus 4.8 (1M context) Co-Authored-By: Claude Fable 5 --- python/sglang/srt/entrypoints/engine.py | 5 +---- .../sglang/srt/model_executor/model_runner.py | 1 + .../remote_instance_weight_transporter.py | 14 ++++++++---- python/sglang/srt/ray/engine.py | 22 ++++++++++++++++++- python/sglang/srt/ray/scheduler_actor.py | 20 ++++++++++++++++- python/sglang/srt/server_args.py | 22 +++++++++++++++++++ 6 files changed, 74 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 25a2a6213032..7387b32bb12b 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -802,10 +802,7 @@ def _launch_subprocesses( # Start the engine info bootstrap server if per-rank info is needed. engine_info_bootstrap_server = None - if ( - server_args.remote_instance_weight_loader_start_seed_via_transfer_engine - and server_args.node_rank == 0 - ): + if server_args.needs_engine_info_bootstrap() and server_args.node_rank == 0: bootstrap_port = server_args.engine_info_bootstrap_port if not is_port_available(bootstrap_port): raise RuntimeError( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 6b32cb645fca..45b0f821003b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -610,6 +610,7 @@ def init_memory_saver_adapter(self): def maybe_init_remote_instance_transfer_engine(self): if self.server_args.remote_instance_weight_loader_use_transfer_engine(): self.remote_instance_weight_transporter.init_engine() + self.remote_instance_weight_transporter.maybe_init_parallelism_config() def maybe_init_expert_location_metadata(self): if self.is_draft_worker: diff --git a/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py b/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py index 476e8070c6d9..8b47f5bc7c9d 100644 --- a/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py +++ b/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py @@ -53,9 +53,15 @@ def init_engine(self): self.session_id = NetworkAddress( local_ip, self.engine.get_rpc_port() ).to_host_port_str() - self.parallelism_config = RankParallelismConfig.from_parallel_state( - self.tp_rank - ) + + def maybe_init_parallelism_config(self) -> None: + # Compute the per-rank parallelism config whenever it will be published to + # the bootstrap server (transfer engine, or RDT/NIXL which needs it without + # one). + if self.server_args.registers_parallelism_config(): + self.parallelism_config = RankParallelismConfig.from_parallel_state( + self.tp_rank + ) def maybe_register_and_publish_weight_info(self) -> None: if ( @@ -75,7 +81,7 @@ def maybe_register_and_publish_weight_info(self) -> None: # The P2P weight-update client needs each rank's parallelism layout to # map training-side parameters onto this rank's shards. if ( - self.server_args.remote_instance_weight_loader_use_transfer_engine() + self.server_args.registers_parallelism_config() and self.parallelism_config is not None ): self._register_parallelism_config_to_bootstrap() diff --git a/python/sglang/srt/ray/engine.py b/python/sglang/srt/ray/engine.py index 57c1d16b86df..b7edffce7f02 100644 --- a/python/sglang/srt/ray/engine.py +++ b/python/sglang/srt/ray/engine.py @@ -203,11 +203,31 @@ def _create_scheduler_actor( return SchedulerActor.options( num_cpus=0, num_gpus=1, + # run_event_loop() blocks its worker thread for the actor's lifetime; allow + # extra threads so a concurrent pull_weights() (RDT weight sync) is not + # starved. Safe because generation is paused during the pull. + max_concurrency=4, + # The http `port` is unique per engine and known to the trainer, so it lets + # RDT discover this engine's scheduler actors by name (ray list_named_actors) + # without an HTTP round-trip, even when several engines share one node. name=( f"sglang_scheduler_node{rank0_node_ip}" f"_dp{dp_rank}_pp{pp_rank}_tp{tp_rank}" - f"_pg{pg.id.hex()[:8]}_bundle{bundle_idx}" + f"_port{server_args.port}_pg{pg.id.hex()[:8]}_bundle{bundle_idx}" ), + # Detached so the trainer (a different Ray driver/job) can discover these via + # list_named_actors / get_actor for RDT pull_weights. Non-detached named + # actors are not listed cross-job. RayEngine.shutdown ray.kills them, so they + # don't leak past the engine's lifetime. + lifetime="detached", + # scheduler_actor uses the absolute GPU id from get_accelerator_ids(); that + # requires Ray NOT to remap CUDA_VISIBLE_DEVICES (else set_device(absolute) + # -> invalid device ordinal). The trainer job sets this in its runtime_env, + # but these actors are created by the sglang subprocess's job, so set it on + # the actor directly. + runtime_env={ + "env_vars": {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"} + }, scheduling_strategy=PlacementGroupSchedulingStrategy( placement_group=pg, placement_group_bundle_index=bundle_idx, diff --git a/python/sglang/srt/ray/scheduler_actor.py b/python/sglang/srt/ray/scheduler_actor.py index e9090ec9a576..de6b4428789f 100644 --- a/python/sglang/srt/ray/scheduler_actor.py +++ b/python/sglang/srt/ray/scheduler_actor.py @@ -16,9 +16,10 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional import ray +from ray import ObjectRef if TYPE_CHECKING: from sglang.srt.server_args import PortArgs, ServerArgs @@ -120,6 +121,23 @@ def get_info(self) -> Dict[str, Any]: """Return scheduler initialization info for handshake.""" return self.scheduler.get_init_info() + def pull_weights(self, weights_refs: List[ObjectRef], param_names: list) -> bool: + """Pull pre-sharded weight bucket from trainer via RDT zero-copy. + + Uses set_target_for_ref to RDMA directly into param.data buffers, + eliminating intermediate receive buffers and copy operations. + + Have to pass weights_refs as a list to avoid resolving upon calling `pull_weights` + """ + from ray.experimental import set_target_for_ref + + model = self.scheduler.tp_worker.model_runner.model + params_dict = dict(model.named_parameters()) + target_buffers = [params_dict[name].data for name in param_names] + set_target_for_ref(weights_refs[0], target_buffers) + ray.get(weights_refs[0]) + return True + def run_event_loop(self) -> None: """Run the scheduler's event loop. Blocks until shutdown.""" try: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 04ba960b9ad2..20bfc104b635 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2669,6 +2669,10 @@ class ServerArgs: bool, "Start seed server via transfer engine backend for remote instance weight loader.", ] = False + enable_engine_info_bootstrap: A[ + bool, + "Start the EngineInfoBootstrapServer and register per-rank parallelism config WITHOUT the mooncake/verbs P2P transfer-engine seeding. Used by RDT (NIXL) weight sync, which needs /parallelism_config but not P2P memory registration.", + ] = False engine_info_bootstrap_port: A[ int, "Port for the engine info bootstrap server. Default is 6789. Must be set explicitly when running multiple instances on the same node.", @@ -8091,6 +8095,24 @@ def remote_instance_weight_loader_use_transfer_engine(self): else: return False + def needs_engine_info_bootstrap(self) -> bool: + """Host the EngineInfoBootstrapServer on this node (rank 0). True for the + transfer-engine seed, and for RDT/NIXL weight sync (which needs the server + to expose parallelism config to the trainer, but not the P2P memory seeding).""" + return ( + self.remote_instance_weight_loader_start_seed_via_transfer_engine + or self.enable_engine_info_bootstrap + ) + + def registers_parallelism_config(self) -> bool: + """Publish this rank's parallelism config to the bootstrap server (hosted + locally by the seed, or remotely). True whenever a transfer engine is in use, + and for RDT/NIXL weight sync.""" + return ( + self.remote_instance_weight_loader_use_transfer_engine() + or self.enable_engine_info_bootstrap + ) + def describe_kv_events_publisher(self) -> Optional[dict]: """Return a structured description of this server's KV-event publisher, or `None` if publishing is disabled / misconfigured. From c6737f6c4da2ef44b8da4f0b5a89d507a25a8c1a Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Mon, 3 Aug 2026 22:33:06 -0700 Subject: [PATCH 02/10] [sglang-miles] gate RDT scheduler-actor options behind --enable-rdt-weight-sync Detached lifetime and threaded actors were applied to every Ray-mode engine. Add a dedicated flag that implies --enable-engine-info-bootstrap and gate them on it. Bind the device in pull_weights, which runs off the event-loop thread. --- python/sglang/srt/ray/engine.py | 12 ++++++------ python/sglang/srt/ray/scheduler_actor.py | 16 +++++++++------- python/sglang/srt/server_args.py | 14 ++++++++++++++ 3 files changed, 29 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/ray/engine.py b/python/sglang/srt/ray/engine.py index b7edffce7f02..a7ade8b8a63e 100644 --- a/python/sglang/srt/ray/engine.py +++ b/python/sglang/srt/ray/engine.py @@ -200,13 +200,15 @@ def _create_scheduler_actor( server_args, tp_rank ) + rdt = server_args.enable_rdt_weight_sync + return SchedulerActor.options( num_cpus=0, num_gpus=1, # run_event_loop() blocks its worker thread for the actor's lifetime; allow - # extra threads so a concurrent pull_weights() (RDT weight sync) is not + # one extra thread so a concurrent pull_weights() (RDT weight sync) is not # starved. Safe because generation is paused during the pull. - max_concurrency=4, + max_concurrency=2 if rdt else 1, # The http `port` is unique per engine and known to the trainer, so it lets # RDT discover this engine's scheduler actors by name (ray list_named_actors) # without an HTTP round-trip, even when several engines share one node. @@ -219,15 +221,13 @@ def _create_scheduler_actor( # list_named_actors / get_actor for RDT pull_weights. Non-detached named # actors are not listed cross-job. RayEngine.shutdown ray.kills them, so they # don't leak past the engine's lifetime. - lifetime="detached", + lifetime="detached" if rdt else None, # scheduler_actor uses the absolute GPU id from get_accelerator_ids(); that # requires Ray NOT to remap CUDA_VISIBLE_DEVICES (else set_device(absolute) # -> invalid device ordinal). The trainer job sets this in its runtime_env, # but these actors are created by the sglang subprocess's job, so set it on # the actor directly. - runtime_env={ - "env_vars": {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"} - }, + runtime_env={"env_vars": {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"}}, scheduling_strategy=PlacementGroupSchedulingStrategy( placement_group=pg, placement_group_bundle_index=bundle_idx, diff --git a/python/sglang/srt/ray/scheduler_actor.py b/python/sglang/srt/ray/scheduler_actor.py index de6b4428789f..d22f124a7883 100644 --- a/python/sglang/srt/ray/scheduler_actor.py +++ b/python/sglang/srt/ray/scheduler_actor.py @@ -121,22 +121,24 @@ def get_info(self) -> Dict[str, Any]: """Return scheduler initialization info for handshake.""" return self.scheduler.get_init_info() - def pull_weights(self, weights_refs: List[ObjectRef], param_names: list) -> bool: - """Pull pre-sharded weight bucket from trainer via RDT zero-copy. + def pull_weights( + self, weights_refs: List[ObjectRef], param_names: List[str] + ) -> None: + """Pull a pre-sharded weight bucket from the trainer via RDT zero-copy. - Uses set_target_for_ref to RDMA directly into param.data buffers, - eliminating intermediate receive buffers and copy operations. - - Have to pass weights_refs as a list to avoid resolving upon calling `pull_weights` + ``weights_refs`` is a list so Ray does not resolve the ref on call. """ + import torch from ray.experimental import set_target_for_ref + # Runs on a different thread than run_event_loop, which owns the device binding. + torch.cuda.set_device(self.scheduler.ps.gpu_id) + model = self.scheduler.tp_worker.model_runner.model params_dict = dict(model.named_parameters()) target_buffers = [params_dict[name].data for name in param_names] set_target_for_ref(weights_refs[0], target_buffers) ray.get(weights_refs[0]) - return True def run_event_loop(self) -> None: """Run the scheduler's event loop. Blocks until shutdown.""" diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 20bfc104b635..373655ea8e04 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2673,6 +2673,10 @@ class ServerArgs: bool, "Start the EngineInfoBootstrapServer and register per-rank parallelism config WITHOUT the mooncake/verbs P2P transfer-engine seeding. Used by RDT (NIXL) weight sync, which needs /parallelism_config but not P2P memory registration.", ] = False + enable_rdt_weight_sync: A[ + bool, + "Expose SchedulerActor.pull_weights for RDT (Ray Direct Transport / NIXL) weight sync from an external trainer job. Requires --use-ray; implies --enable-engine-info-bootstrap.", + ] = False engine_info_bootstrap_port: A[ int, "Port for the engine info bootstrap server. Default is 6789. Must be set explicitly when running multiple instances on the same node.", @@ -2931,6 +2935,8 @@ def __post_init__(self): # _handle_model_specific_adjustments never runs. self._resolved_overrides = [] + self._handle_rdt_weight_sync() + if self.model_path.lower() in ["none", "dummy"]: return @@ -3164,6 +3170,14 @@ def _handle_model_source_paths(self): ): ObjectStorageModel.download_and_get_path(self.tokenizer_path) + def _handle_rdt_weight_sync(self): + if not self.enable_rdt_weight_sync: + return + assert ( + self.use_ray + ), "--enable-rdt-weight-sync requires --use-ray: the trainer pulls weights through named SchedulerActors." + self.enable_engine_info_bootstrap = True + def _handle_pd_disaggregation(self): from sglang.srt.arg_groups.pd_disaggregation_hook import ( handle_pd_disaggregation, From 7c15e1f16206c571c09169890f072a24673c5549 Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Mon, 3 Aug 2026 22:34:27 -0700 Subject: [PATCH 03/10] Trim RDT comments to sglang core style --- .../remote_instance_weight_transporter.py | 3 --- python/sglang/srt/ray/engine.py | 23 +++++++------------ python/sglang/srt/server_args.py | 10 +++----- 3 files changed, 11 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py b/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py index 8b47f5bc7c9d..ba85474997e5 100644 --- a/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py +++ b/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py @@ -55,9 +55,6 @@ def init_engine(self): ).to_host_port_str() def maybe_init_parallelism_config(self) -> None: - # Compute the per-rank parallelism config whenever it will be published to - # the bootstrap server (transfer engine, or RDT/NIXL which needs it without - # one). if self.server_args.registers_parallelism_config(): self.parallelism_config = RankParallelismConfig.from_parallel_state( self.tp_rank diff --git a/python/sglang/srt/ray/engine.py b/python/sglang/srt/ray/engine.py index a7ade8b8a63e..5d2562416d88 100644 --- a/python/sglang/srt/ray/engine.py +++ b/python/sglang/srt/ray/engine.py @@ -205,28 +205,21 @@ def _create_scheduler_actor( return SchedulerActor.options( num_cpus=0, num_gpus=1, - # run_event_loop() blocks its worker thread for the actor's lifetime; allow - # one extra thread so a concurrent pull_weights() (RDT weight sync) is not - # starved. Safe because generation is paused during the pull. + # run_event_loop() blocks one thread for the actor's lifetime; leave a spare + # for pull_weights, which the trainer calls while generation is paused. max_concurrency=2 if rdt else 1, - # The http `port` is unique per engine and known to the trainer, so it lets - # RDT discover this engine's scheduler actors by name (ray list_named_actors) - # without an HTTP round-trip, even when several engines share one node. + # The http `port` disambiguates engines co-located on one node, letting the + # trainer find these actors via list_named_actors. name=( f"sglang_scheduler_node{rank0_node_ip}" f"_dp{dp_rank}_pp{pp_rank}_tp{tp_rank}" f"_port{server_args.port}_pg{pg.id.hex()[:8]}_bundle{bundle_idx}" ), - # Detached so the trainer (a different Ray driver/job) can discover these via - # list_named_actors / get_actor for RDT pull_weights. Non-detached named - # actors are not listed cross-job. RayEngine.shutdown ray.kills them, so they - # don't leak past the engine's lifetime. + # Non-detached named actors are not listed cross-job, so the trainer (a + # separate Ray job) could not discover them. RayEngine.shutdown kills these. lifetime="detached" if rdt else None, - # scheduler_actor uses the absolute GPU id from get_accelerator_ids(); that - # requires Ray NOT to remap CUDA_VISIBLE_DEVICES (else set_device(absolute) - # -> invalid device ordinal). The trainer job sets this in its runtime_env, - # but these actors are created by the sglang subprocess's job, so set it on - # the actor directly. + # SchedulerActor calls set_device() with the absolute id from + # get_accelerator_ids(), which is only valid if Ray leaves the mask alone. runtime_env={"env_vars": {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"}}, scheduling_strategy=PlacementGroupSchedulingStrategy( placement_group=pg, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 373655ea8e04..d0c7288e3fd0 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2671,7 +2671,7 @@ class ServerArgs: ] = False enable_engine_info_bootstrap: A[ bool, - "Start the EngineInfoBootstrapServer and register per-rank parallelism config WITHOUT the mooncake/verbs P2P transfer-engine seeding. Used by RDT (NIXL) weight sync, which needs /parallelism_config but not P2P memory registration.", + "Start the EngineInfoBootstrapServer and register per-rank parallelism config, without the mooncake/verbs P2P transfer-engine seeding.", ] = False enable_rdt_weight_sync: A[ bool, @@ -8110,18 +8110,14 @@ def remote_instance_weight_loader_use_transfer_engine(self): return False def needs_engine_info_bootstrap(self) -> bool: - """Host the EngineInfoBootstrapServer on this node (rank 0). True for the - transfer-engine seed, and for RDT/NIXL weight sync (which needs the server - to expose parallelism config to the trainer, but not the P2P memory seeding).""" + """Whether this node (rank 0) hosts the EngineInfoBootstrapServer.""" return ( self.remote_instance_weight_loader_start_seed_via_transfer_engine or self.enable_engine_info_bootstrap ) def registers_parallelism_config(self) -> bool: - """Publish this rank's parallelism config to the bootstrap server (hosted - locally by the seed, or remotely). True whenever a transfer engine is in use, - and for RDT/NIXL weight sync.""" + """Whether this rank publishes its parallelism config to the bootstrap server.""" return ( self.remote_instance_weight_loader_use_transfer_engine() or self.enable_engine_info_bootstrap From ce4c918949e5d8138caf763578d757c5e4766436 Mon Sep 17 00:00:00 2001 From: xyuzh Date: Mon, 10 Aug 2026 16:12:53 -0700 Subject: [PATCH 04/10] [sglang-miles] accept a caller-supplied placement group for the Ray backend launch_server() unconditionally cleared server_args.placement_group, so a Ray-backend server launched from a parent job could not join an existing placement group and would auto-create a second one, double-booking the rollout GPUs. Honor a caller-supplied placement_group instead, and connect to the running cluster with the caller's runtime env when the launching process has no Ray context of its own (the mp.Process child loses it). Add SchedulerActor.register_weight_for_rdt() so the destination pins its model parameters with NIXL once, rather than re-pinning and re-handshaking on every RDT flush. Skipped under --enable-memory-saver, where parameter storages are not stably resident. Bump the ray extra to >=2.56.0 for ray.experimental.register_nixl_memory. --- python/pyproject.toml | 2 +- python/sglang/srt/ray/http_server.py | 13 ++++++++++++- python/sglang/srt/ray/scheduler_actor.py | 13 +++++++++++++ 3 files changed, 26 insertions(+), 2 deletions(-) diff --git a/python/pyproject.toml b/python/pyproject.toml index b934e83351f8..1246b2b00420 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -123,7 +123,7 @@ diffusion = [ ] ray = [ - "ray[default]>=2.55.1", + "ray[default]>=2.56.0", ] tracing = [ diff --git a/python/sglang/srt/ray/http_server.py b/python/sglang/srt/ray/http_server.py index 7d7e23567aa2..ac955f614796 100644 --- a/python/sglang/srt/ray/http_server.py +++ b/python/sglang/srt/ray/http_server.py @@ -15,6 +15,8 @@ from typing import Callable, Optional +import ray + from sglang.srt.entrypoints.engine import ( init_tokenizer_manager, run_detokenizer_process, @@ -44,7 +46,16 @@ def launch_server( if execute_warmup_func is None: execute_warmup_func = _execute_server_warmup - server_args.override("ray.http_server.clear_placement_group", placement_group=None) + placement_group = getattr(server_args, "placement_group", None) + if placement_group is not None and not ray.is_initialized(): + ray.init( + address="auto", + runtime_env=getattr(server_args, "ray_runtime_env", None), + ) + server_args.override( + "ray.http_server.placement_group", + placement_group=placement_group, + ) ( tokenizer_manager, diff --git a/python/sglang/srt/ray/scheduler_actor.py b/python/sglang/srt/ray/scheduler_actor.py index d22f124a7883..e3a95b704a40 100644 --- a/python/sglang/srt/ray/scheduler_actor.py +++ b/python/sglang/srt/ray/scheduler_actor.py @@ -121,6 +121,19 @@ def get_info(self) -> Dict[str, Any]: """Return scheduler initialization info for handshake.""" return self.scheduler.get_init_info() + def register_weight_for_rdt(self) -> None: + """Pin model parameters with NIXL for repeated RDT pulls.""" + if self.scheduler.server_args.enable_memory_saver: + return + + import torch + from ray.experimental import register_nixl_memory + + torch.cuda.set_device(self.scheduler.ps.gpu_id) + model = self.scheduler.tp_worker.model_runner.model + for _, param in model.named_parameters(): + register_nixl_memory(param.data) + def pull_weights( self, weights_refs: List[ObjectRef], param_names: List[str] ) -> None: From c75459ca779fdb205da3bab525f8b9162c95fe94 Mon Sep 17 00:00:00 2001 From: xyuzh Date: Wed, 12 Aug 2026 12:06:17 -0700 Subject: [PATCH 05/10] [sglang-miles] export get_scheduler_actor_name so callers can find the actors The SchedulerActor name format was inlined in _create_scheduler_actor, so an external trainer holding the placement bundles could only find the actors by scanning list_named_actors for substring matches. Move the format into a get_scheduler_actor_name() helper and re-export it from sglang.srt.ray, so a caller that knows (node ip, dp/pp/tp rank, port, bundle index) can rebuild the name and ray.get_actor() it directly. The placement-group hex is dropped from the name: it is not knowable from the caller side, and the http port already disambiguates engines co-located on one node. --- python/sglang/srt/ray/__init__.py | 4 ++-- python/sglang/srt/ray/engine.py | 30 ++++++++++++++++++++++++------ 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/ray/__init__.py b/python/sglang/srt/ray/__init__.py index 5927c789f926..57a081914468 100644 --- a/python/sglang/srt/ray/__init__.py +++ b/python/sglang/srt/ray/__init__.py @@ -1,3 +1,3 @@ -from sglang.srt.ray.engine import RayEngine +from sglang.srt.ray.engine import RayEngine, get_scheduler_actor_name -__all__ = ["RayEngine"] +__all__ = ["RayEngine", "get_scheduler_actor_name"] diff --git a/python/sglang/srt/ray/engine.py b/python/sglang/srt/ray/engine.py index 5d2562416d88..7094f858901e 100644 --- a/python/sglang/srt/ray/engine.py +++ b/python/sglang/srt/ray/engine.py @@ -175,6 +175,23 @@ def _validate_custom_placement_group(pg: PlacementGroup, world_size: int) -> Non ) +def get_scheduler_actor_name( + *, + rank0_node_ip: str, + dp_rank: int, + pp_rank: int, + tp_rank: int, + port: int, + bundle_idx: int, +) -> str: + """Return the Ray actor name for a SchedulerActor. Can be used to retrive scheduler ray actors""" + return ( + f"sglang_scheduler_node{rank0_node_ip}" + f"_dp{dp_rank}_pp{pp_rank}_tp{tp_rank}" + f"_port{port}_bundle{bundle_idx}" + ) + + def _create_scheduler_actor( pg: PlacementGroup, bundle_idx: int, @@ -208,12 +225,13 @@ def _create_scheduler_actor( # run_event_loop() blocks one thread for the actor's lifetime; leave a spare # for pull_weights, which the trainer calls while generation is paused. max_concurrency=2 if rdt else 1, - # The http `port` disambiguates engines co-located on one node, letting the - # trainer find these actors via list_named_actors. - name=( - f"sglang_scheduler_node{rank0_node_ip}" - f"_dp{dp_rank}_pp{pp_rank}_tp{tp_rank}" - f"_port{server_args.port}_pg{pg.id.hex()[:8]}_bundle{bundle_idx}" + name=get_scheduler_actor_name( + rank0_node_ip=rank0_node_ip, + dp_rank=dp_rank, + pp_rank=pp_rank, + tp_rank=tp_rank, + port=server_args.port, + bundle_idx=bundle_idx, ), # Non-detached named actors are not listed cross-job, so the trainer (a # separate Ray job) could not discover them. RayEngine.shutdown kills these. From d8ecc7ce4697e61488ee29154650d00567834d6e Mon Sep 17 00:00:00 2001 From: xyuzh Date: Thu, 13 Aug 2026 10:20:01 -0700 Subject: [PATCH 06/10] [sglang-miles] join the caller's Ray namespace when connecting to the cluster The Ray-backend launch path connects with ray.init(address="auto"), which lands in a fresh anonymous namespace when none is given. SchedulerActors then register there instead of in the namespace the rest of the deployment uses, so discovery by name from another job has to scan every namespace. Honor a caller-supplied ray_namespace, alongside the placement group and runtime env already threaded through server_args. --- python/sglang/srt/ray/http_server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/sglang/srt/ray/http_server.py b/python/sglang/srt/ray/http_server.py index ac955f614796..4b9d1510f7c7 100644 --- a/python/sglang/srt/ray/http_server.py +++ b/python/sglang/srt/ray/http_server.py @@ -51,6 +51,7 @@ def launch_server( ray.init( address="auto", runtime_env=getattr(server_args, "ray_runtime_env", None), + namespace=getattr(server_args, "ray_namespace", None), ) server_args.override( "ray.http_server.placement_group", From 2ebffd34b2ea47ab54dd9a0bc8fbb775b7eb3fb3 Mon Sep 17 00:00:00 2001 From: xyuzh Date: Tue, 18 Aug 2026 15:37:26 -0700 Subject: [PATCH 07/10] [sglang-miles] split launch_server into launch_engine + serve_http Callers hosting the HTTP server in their own Ray actor can now call launch_engine to get the SchedulerActor handles back, then run serve_http on a thread. With handles returned directly, the schedulers no longer need detached lifetimes or name-based discovery, so the actor names key on the placement group ID instead of the server port. --- python/sglang/srt/ray/__init__.py | 5 +- python/sglang/srt/ray/engine.py | 10 ++-- python/sglang/srt/ray/http_server.py | 77 ++++++++++++++++++---------- python/sglang/srt/server_args.py | 4 +- 4 files changed, 59 insertions(+), 37 deletions(-) diff --git a/python/sglang/srt/ray/__init__.py b/python/sglang/srt/ray/__init__.py index 57a081914468..ba619e3bc7a5 100644 --- a/python/sglang/srt/ray/__init__.py +++ b/python/sglang/srt/ray/__init__.py @@ -1,3 +1,4 @@ -from sglang.srt.ray.engine import RayEngine, get_scheduler_actor_name +from sglang.srt.ray.engine import RayEngine +from sglang.srt.ray.http_server import launch_engine, launch_server, serve_http -__all__ = ["RayEngine", "get_scheduler_actor_name"] +__all__ = ["RayEngine", "launch_engine", "launch_server", "serve_http"] diff --git a/python/sglang/srt/ray/engine.py b/python/sglang/srt/ray/engine.py index 7094f858901e..ff2c13f0cca9 100644 --- a/python/sglang/srt/ray/engine.py +++ b/python/sglang/srt/ray/engine.py @@ -181,14 +181,13 @@ def get_scheduler_actor_name( dp_rank: int, pp_rank: int, tp_rank: int, - port: int, + pg_id_hex: str, bundle_idx: int, ) -> str: - """Return the Ray actor name for a SchedulerActor. Can be used to retrive scheduler ray actors""" return ( f"sglang_scheduler_node{rank0_node_ip}" f"_dp{dp_rank}_pp{pp_rank}_tp{tp_rank}" - f"_port{port}_bundle{bundle_idx}" + f"_pg{pg_id_hex}_bundle{bundle_idx}" ) @@ -230,12 +229,9 @@ def _create_scheduler_actor( dp_rank=dp_rank, pp_rank=pp_rank, tp_rank=tp_rank, - port=server_args.port, + pg_id_hex=pg.id.hex()[:8], bundle_idx=bundle_idx, ), - # Non-detached named actors are not listed cross-job, so the trainer (a - # separate Ray job) could not discover them. RayEngine.shutdown kills these. - lifetime="detached" if rdt else None, # SchedulerActor calls set_device() with the absolute id from # get_accelerator_ids(), which is only valid if Ray leaves the mask alone. runtime_env={"env_vars": {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"}}, diff --git a/python/sglang/srt/ray/http_server.py b/python/sglang/srt/ray/http_server.py index 4b9d1510f7c7..17a67f56880a 100644 --- a/python/sglang/srt/ray/http_server.py +++ b/python/sglang/srt/ray/http_server.py @@ -15,8 +15,6 @@ from typing import Callable, Optional -import ray - from sglang.srt.entrypoints.engine import ( init_tokenizer_manager, run_detokenizer_process, @@ -25,51 +23,53 @@ from sglang.srt.server_args import ServerArgs -def launch_server( +def launch_engine( server_args: ServerArgs, init_tokenizer_manager_func: Callable = init_tokenizer_manager, run_scheduler_process_func: Callable = run_scheduler_process, run_detokenizer_process_func: Callable = run_detokenizer_process, - execute_warmup_func: Optional[Callable] = None, - launch_callback: Optional[Callable[[], None]] = None, ): - """Launch HTTP server with Ray-based scheduler actors. + """Create RayEngine subprocesses / SchedulerActors. Does not start HTTP. - Mirrors http_server.launch_server() but uses RayEngine for scheduler launching. + Returns the ``_launch_subprocesses`` 5-tuple: + ``(tokenizer_manager, template_manager, port_args, scheduler_init_result, + subprocess_watchdog)``. ``scheduler_init_result.scheduler_actors`` are the + SchedulerActor handles. """ + from sglang.srt.ray.engine import RayEngine + + server_args.override("ray.http_server.clear_placement_group", placement_group=None) + + return RayEngine._launch_subprocesses( + server_args, + init_tokenizer_manager_func=init_tokenizer_manager_func, + run_scheduler_process_func=run_scheduler_process_func, + run_detokenizer_process_func=run_detokenizer_process_func, + ) + + +def serve_http( + engine, + server_args: ServerArgs, + execute_warmup_func: Optional[Callable] = None, + launch_callback: Optional[Callable[[], None]] = None, +): + """Block in uvicorn. ``engine`` is the 5-tuple from ``launch_engine``.""" from sglang.srt.entrypoints.http_server import ( _execute_server_warmup, _setup_and_run_http_server, ) - from sglang.srt.ray.engine import RayEngine if execute_warmup_func is None: execute_warmup_func = _execute_server_warmup - placement_group = getattr(server_args, "placement_group", None) - if placement_group is not None and not ray.is_initialized(): - ray.init( - address="auto", - runtime_env=getattr(server_args, "ray_runtime_env", None), - namespace=getattr(server_args, "ray_namespace", None), - ) - server_args.override( - "ray.http_server.placement_group", - placement_group=placement_group, - ) - ( tokenizer_manager, template_manager, port_args, scheduler_init_result, subprocess_watchdog, - ) = RayEngine._launch_subprocesses( - server_args, - init_tokenizer_manager_func=init_tokenizer_manager_func, - run_scheduler_process_func=run_scheduler_process_func, - run_detokenizer_process_func=run_detokenizer_process_func, - ) + ) = engine _setup_and_run_http_server( server_args, @@ -81,3 +81,28 @@ def launch_server( execute_warmup_func=execute_warmup_func, launch_callback=launch_callback, ) + + +def launch_server( + server_args: ServerArgs, + init_tokenizer_manager_func: Callable = init_tokenizer_manager, + run_scheduler_process_func: Callable = run_scheduler_process, + run_detokenizer_process_func: Callable = run_detokenizer_process, + execute_warmup_func: Optional[Callable] = None, + launch_callback: Optional[Callable[[], None]] = None, +): + """Launch HTTP server with Ray-based scheduler actors. + + Mirrors http_server.launch_server() but uses RayEngine for scheduler launching. + """ + serve_http( + launch_engine( + server_args, + init_tokenizer_manager_func=init_tokenizer_manager_func, + run_scheduler_process_func=run_scheduler_process_func, + run_detokenizer_process_func=run_detokenizer_process_func, + ), + server_args, + execute_warmup_func=execute_warmup_func, + launch_callback=launch_callback, + ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d0c7288e3fd0..8687e29a4d5c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2675,7 +2675,7 @@ class ServerArgs: ] = False enable_rdt_weight_sync: A[ bool, - "Expose SchedulerActor.pull_weights for RDT (Ray Direct Transport / NIXL) weight sync from an external trainer job. Requires --use-ray; implies --enable-engine-info-bootstrap.", + "Expose SchedulerActor.pull_weights for RDT (Ray Direct Transport / NIXL) weight sync. Requires --use-ray; implies --enable-engine-info-bootstrap.", ] = False engine_info_bootstrap_port: A[ int, @@ -3175,7 +3175,7 @@ def _handle_rdt_weight_sync(self): return assert ( self.use_ray - ), "--enable-rdt-weight-sync requires --use-ray: the trainer pulls weights through named SchedulerActors." + ), "--enable-rdt-weight-sync requires --use-ray: the trainer pulls weights through SchedulerActors." self.enable_engine_info_bootstrap = True def _handle_pd_disaggregation(self): From 8e017dda4709355d07c7c3e22316afb514088cc6 Mon Sep 17 00:00:00 2001 From: xyuzh Date: Wed, 19 Aug 2026 09:51:45 -0700 Subject: [PATCH 08/10] [sglang-miles] keep the caller-supplied placement group in launch_engine launch_engine cleared server_args.placement_group before launching, so SchedulerActors fell back to their own scheduling instead of the group the caller reserved. Honor the override the caller set. --- python/sglang/srt/ray/http_server.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/python/sglang/srt/ray/http_server.py b/python/sglang/srt/ray/http_server.py index 17a67f56880a..909e4b8b3eeb 100644 --- a/python/sglang/srt/ray/http_server.py +++ b/python/sglang/srt/ray/http_server.py @@ -29,17 +29,9 @@ def launch_engine( run_scheduler_process_func: Callable = run_scheduler_process, run_detokenizer_process_func: Callable = run_detokenizer_process, ): - """Create RayEngine subprocesses / SchedulerActors. Does not start HTTP. - - Returns the ``_launch_subprocesses`` 5-tuple: - ``(tokenizer_manager, template_manager, port_args, scheduler_init_result, - subprocess_watchdog)``. ``scheduler_init_result.scheduler_actors`` are the - SchedulerActor handles. - """ + """Create RayEngine subprocesses / SchedulerActors.""" from sglang.srt.ray.engine import RayEngine - server_args.override("ray.http_server.clear_placement_group", placement_group=None) - return RayEngine._launch_subprocesses( server_args, init_tokenizer_manager_func=init_tokenizer_manager_func, From 1fcbc8b4f496b9a7936d6273b513fcdceda141e8 Mon Sep 17 00:00:00 2001 From: xyuzh Date: Wed, 19 Aug 2026 11:22:23 -0700 Subject: [PATCH 09/10] [sglang-miles] take the caller placement group as an explicit argument launch_engine and RayEngine now accept placement_group directly and carry it to _launch_subprocesses through a contextvar, replacing the dynamic server_args.placement_group attribute that had to be manually re-attached after every dataclasses.replace. Pairs with the separate PR that has RayDataParallelController accept is_custom_pg, which this passes. --- python/sglang/srt/ray/engine.py | 53 +++++++++++++++++++--------- python/sglang/srt/ray/http_server.py | 19 ++++++---- 2 files changed, 49 insertions(+), 23 deletions(-) diff --git a/python/sglang/srt/ray/engine.py b/python/sglang/srt/ray/engine.py index ff2c13f0cca9..9f6cacefd4d0 100644 --- a/python/sglang/srt/ray/engine.py +++ b/python/sglang/srt/ray/engine.py @@ -15,9 +15,11 @@ from __future__ import annotations +import contextvars import dataclasses import logging import threading +from contextlib import contextmanager from typing import Callable, List, Optional import ray @@ -36,6 +38,20 @@ logger = logging.getLogger(__name__) +_caller_placement_group: contextvars.ContextVar[Optional[PlacementGroup]] = ( + contextvars.ContextVar("sglang_ray_caller_placement_group", default=None) +) + + +@contextmanager +def _placement_group_context(placement_group: Optional[PlacementGroup]): + """Expose launch-only state while Engine synchronously starts schedulers.""" + token = _caller_placement_group.set(placement_group) + try: + yield + finally: + _caller_placement_group.reset(token) + @dataclasses.dataclass class RaySchedulerInitResult(SchedulerInitResult): @@ -254,15 +270,17 @@ def _create_scheduler_actor( class RayEngine(Engine): - """Engine using Ray actors for scheduler processes.""" + """Engine using Ray actors for scheduler processes. + + Same constructor kwargs as :class:`Engine`, plus ``placement_group`` (a Ray + PlacementGroup handle, not a ServerArgs field). + """ - def __init__(self, **kwargs): - placement_group = kwargs.pop("placement_group", None) - if "log_level" not in kwargs: - kwargs["log_level"] = "error" - server_args = ServerArgs(**kwargs) - server_args.override("ray.placement_group", placement_group=placement_group) - super().__init__(server_args=server_args) + def __init__( + self, *, placement_group: Optional[PlacementGroup] = None, **kwargs + ): + with _placement_group_context(placement_group): + super().__init__(**kwargs) def shutdown(self): """Shutdown the engine — kill Ray scheduler actors then local processes.""" @@ -286,7 +304,8 @@ def _launch_scheduler_processes( Tuple of (RaySchedulerInitResult, None). scheduler_procs is None since Ray uses actors instead of mp.Process. """ - pg = server_args.placement_group or ray.util.get_current_placement_group() + placement_group = _caller_placement_group.get() + pg = placement_group or ray.util.get_current_placement_group() if pg is None: from ray.util.placement_group import ( placement_group as create_placement_group, @@ -315,7 +334,7 @@ def _launch_scheduler_processes( ) ray.get(pg.ready()) - is_custom_pg = server_args.placement_group is not None + is_custom_pg = placement_group is not None nnodes = server_args.nnodes world_size = _compute_world_size(server_args) @@ -451,6 +470,7 @@ def wait_for_completion(): pg, bundle_for_node, rank0_node_ip, + is_custom_pg, ), None, ) @@ -463,6 +483,7 @@ def _launch_dp_scheduler_processes( pg, bundle_for_node: Optional[List[int]], rank0_node_ip: str, + is_custom_pg: bool = False, ) -> RaySchedulerInitResult: """Launch DP schedulers via RayDataParallelController.""" from sglang.srt.ray.data_parallel_controller import ( @@ -488,16 +509,16 @@ def _launch_dp_scheduler_processes( server_args, dist_init_addr=f"{rank0_node_ip}:{port_args.nccl_port}", ) - # dataclasses.replace only copies declared fields; placement_group is - # a dynamic attribute that must be manually appended after the rebuild. - dp_server_args.override( - "ray.placement_group", placement_group=server_args.placement_group - ) # Create the DP controller in-process. This blocks until all actors # are initialized and their event loops have started. controller = RayDataParallelController( - dp_server_args, port_args, pg, bundle_for_node, rank0_node_ip + dp_server_args, + port_args, + pg, + bundle_for_node, + rank0_node_ip, + is_custom_pg, ) # Start the DP controller's event loop in a daemon thread. diff --git a/python/sglang/srt/ray/http_server.py b/python/sglang/srt/ray/http_server.py index 909e4b8b3eeb..f8c568956934 100644 --- a/python/sglang/srt/ray/http_server.py +++ b/python/sglang/srt/ray/http_server.py @@ -15,6 +15,8 @@ from typing import Callable, Optional +from ray.util.placement_group import PlacementGroup + from sglang.srt.entrypoints.engine import ( init_tokenizer_manager, run_detokenizer_process, @@ -28,16 +30,19 @@ def launch_engine( init_tokenizer_manager_func: Callable = init_tokenizer_manager, run_scheduler_process_func: Callable = run_scheduler_process, run_detokenizer_process_func: Callable = run_detokenizer_process, + *, + placement_group: Optional[PlacementGroup] = None, ): """Create RayEngine subprocesses / SchedulerActors.""" - from sglang.srt.ray.engine import RayEngine + from sglang.srt.ray.engine import RayEngine, _placement_group_context - return RayEngine._launch_subprocesses( - server_args, - init_tokenizer_manager_func=init_tokenizer_manager_func, - run_scheduler_process_func=run_scheduler_process_func, - run_detokenizer_process_func=run_detokenizer_process_func, - ) + with _placement_group_context(placement_group): + return RayEngine._launch_subprocesses( + server_args, + init_tokenizer_manager_func=init_tokenizer_manager_func, + run_scheduler_process_func=run_scheduler_process_func, + run_detokenizer_process_func=run_detokenizer_process_func, + ) def serve_http( From a7bb39200d1426fa0f4c797dc07f81cd739dba41 Mon Sep 17 00:00:00 2001 From: xyuzh Date: Thu, 20 Aug 2026 21:36:47 -0700 Subject: [PATCH 10/10] [sglang-miles] appease black on the RayEngine constructor signature --- python/sglang/srt/ray/engine.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/sglang/srt/ray/engine.py b/python/sglang/srt/ray/engine.py index 9f6cacefd4d0..3e3fc400ba0b 100644 --- a/python/sglang/srt/ray/engine.py +++ b/python/sglang/srt/ray/engine.py @@ -276,9 +276,7 @@ class RayEngine(Engine): PlacementGroup handle, not a ServerArgs field). """ - def __init__( - self, *, placement_group: Optional[PlacementGroup] = None, **kwargs - ): + def __init__(self, *, placement_group: Optional[PlacementGroup] = None, **kwargs): with _placement_group_context(placement_group): super().__init__(**kwargs)