sglang 0.5.2rc1__py3-none-any.whl → 0.5.3rc0__py3-none-any.whl

sglang/srt/offloader.py

@@ -38,6 +38,10 @@ class BaseOffloader(ABC):
     def post_init(self):
         pass
 
+    @property
+    def forbid_copy_engine_usage(self):
+        return False
+
 
 class NoopOffloader(BaseOffloader):
     pass
@@ -233,6 +237,10 @@ class OffloaderV2(BaseOffloader):
         for i in range(self.prefetch_step):
             self.offloaders[i].start_onload()
 
+    @property
+    def forbid_copy_engine_usage(self):
+        return self.mode == "cpu"
+
 
 def _hook_module_forward_for_offloader(index, module, offloaders, prefetch_step):
     def _on_forward_end():
@@ -398,14 +406,30 @@ class _ShmCpuParamOffloader(_BaseParamOffloader):
         return self.shm_cpu_data.to("cuda", non_blocking=True)
 
 
+def update_param(param, new_tensor):
+    """Update parameter while keeping properties needed by Offloader (e.g. pinned host memory)."""
+
+    if param.device == new_tensor.device:
+        param.data = new_tensor
+    else:
+        assert param.device == torch.device(
+            "cpu"
+        ), f"{param.device=} {new_tensor.device=}"
+        param.data = _create_cpu_data(new_tensor, pin_memory=True)
+
+
 def _move_param_to_cpu(param, pin_memory: bool):
+    param.data = _create_cpu_data(param.data, pin_memory=pin_memory)
+
+
+def _create_cpu_data(data, pin_memory: bool):
     cpu_data = _empty_strided_like(
-        param.data,
+        data,
         device="cpu",
         pin_memory=pin_memory,
     )
-    cpu_data.copy_(param.data)
-    param.data = cpu_data
+    cpu_data.copy_(data)
+    return cpu_data
 
 
 def _move_param_to_meta(module, param_name):

sglang/srt/{reasoning_parser.py → parser/reasoning_parser.py}

@@ -1,7 +1,7 @@
 import re
 from typing import Dict, Optional, Tuple, Type
 
-from sglang.srt.harmony_parser import HarmonyParser
+from sglang.srt.parser.harmony_parser import HarmonyParser
 
 
 class StreamingParseResult:

sglang/srt/remote_instance_weight_loader_utils.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from typing import List
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+def trigger_init_weights_send_group_for_remote_instance_request(
+    remote_instance_weight_loader_seed_instance_ip: str,
+    remote_instance_weight_loader_seed_instance_service_port: int,
+    remote_instance_weight_loader_send_weights_group_ports: List[int],
+    remote_instance_weight_loader_client_id: str,
+):
+    seed_instance_service_url = f"http://{remote_instance_weight_loader_seed_instance_ip}:{remote_instance_weight_loader_seed_instance_service_port}"
+    # Only support loading weights from instance with same parallelism strategy.
+    # Per TP rank pair between seed and dst instances will build a communication group for sending weights.
+    # i.e. seed TP 0 <-> dst TP 0, seed TP 1 <-> dst TP 1, etc.
+    # Each communication group will have a world size 2.
+    try:
+        requests.post(
+            f"{seed_instance_service_url}/init_weights_send_group_for_remote_instance",
+            json={
+                "master_address": remote_instance_weight_loader_seed_instance_ip,
+                "ports": (
+                    ",".join(
+                        str(p)
+                        for p in remote_instance_weight_loader_send_weights_group_ports
+                    )
+                ),
+                "group_rank": 0,
+                "world_size": 2,
+                "group_name": f"send_weights_{remote_instance_weight_loader_client_id}",
+                "backend": "nccl",
+            },
+        )
+    except Exception as e:
+        logger.error(
+            f"Failed to trigger init_weights_send_group_for_remote_instance_request to seed instance {seed_instance_service_url}: {e}."
+        )
+        raise
+
+
+def trigger_transferring_weights_request(
+    remote_instance_weight_loader_seed_instance_ip: str,
+    remote_instance_weight_loader_seed_instance_service_port: int,
+    remote_instance_weight_loader_send_weights_group_ports: List[int],
+    remote_instance_weight_loader_client_id: str,
+):
+    seed_instance_service_url = f"http://{remote_instance_weight_loader_seed_instance_ip}:{remote_instance_weight_loader_seed_instance_service_port}"
+    try:
+        requests.post(
+            f"{seed_instance_service_url}/send_weights_to_remote_instance",
+            json={
+                "master_address": remote_instance_weight_loader_seed_instance_ip,
+                "ports": (
+                    ",".join(
+                        str(p)
+                        for p in remote_instance_weight_loader_send_weights_group_ports
+                    )
+                ),
+                "group_name": f"send_weights_{remote_instance_weight_loader_client_id}",
+            },
+        )
+    except Exception as e:
+        logger.error(f"Failed to trigger send weights to remote instance request: {e}")
+        raise

sglang/srt/sampling/sampling_batch_info.py

@@ -67,28 +67,31 @@ class SamplingBatchInfo:
     logit_bias: Optional[torch.Tensor] = None
 
     @classmethod
-    def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int):
+    def _get_global_server_args_dict(cls):
         from sglang.srt.managers.schedule_batch import global_server_args_dict
 
+        return global_server_args_dict
+
+    @classmethod
+    def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int):
+        global_server_args_dict = cls._get_global_server_args_dict()
+
         reqs = batch.reqs
         device = batch.device
-        temperatures = (
-            torch.tensor(
-                [r.sampling_params.temperature for r in reqs],
-                dtype=torch.float,
-            )
-            .view(-1, 1)
-            .to(device, non_blocking=True)
-        )
+        temperatures = torch.tensor(
+            [r.sampling_params.temperature for r in reqs],
+            dtype=torch.float,
+            device=device,
+        ).view(-1, 1)
         top_ps = torch.tensor(
-            [r.sampling_params.top_p for r in reqs], dtype=torch.float
-        ).to(device, non_blocking=True)
+            [r.sampling_params.top_p for r in reqs], dtype=torch.float, device=device
+        )
         top_ks = torch.tensor(
-            [r.sampling_params.top_k for r in reqs], dtype=torch.int32
-        ).to(device, non_blocking=True)
+            [r.sampling_params.top_k for r in reqs], dtype=torch.int32, device=device
+        )
         min_ps = torch.tensor(
-            [r.sampling_params.min_p for r in reqs], dtype=torch.float
-        ).to(device, non_blocking=True)
+            [r.sampling_params.min_p for r in reqs], dtype=torch.float, device=device
+        )
 
         logit_bias = None
         if any(r.sampling_params.logit_bias is not None for r in reqs):