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

sglang/srt/debug_utils/dumper.py

@@ -36,6 +36,15 @@ class _Dumper:
         self._forward_pass_id = 0
 
     def on_forward_pass_start(self):
+        """This should be called on all ranks."""
+
+        if not self._enable:
+            return
+
+        # Users may want to `dump` only on some ranks, thus determine name here
+        if self._partial_name is None:
+            self._partial_name = _get_partial_name()
+
         self._forward_pass_id += 1
         print(
             f"[Dumper] [{time.time()}] on_forward_pass_start id={self._forward_pass_id}"
@@ -48,12 +57,10 @@ class _Dumper:
         assert (
             self._forward_pass_id >= 1
         ), "Do you forget to call `dumper.on_forward_pass_start()`?"
+        assert self._partial_name is not None
         self._dump_index += 1
 
-        if self._partial_name is None:
-            self._partial_name = _get_partial_name()
-
-        rank = dist.get_rank()
+        rank = _get_rank()
         full_kwargs = dict(
             forward_pass_id=self._forward_pass_id,
             rank=rank,
@@ -80,12 +87,20 @@ class _Dumper:
 
 
 def _get_partial_name():
-    rank = dist.get_rank()
+    rank = _get_rank()
     object_list = [str(time.time()) if rank == 0 else None]
-    dist.broadcast_object_list(object_list, device="cuda")
+    if dist.is_initialized():
+        dist.broadcast_object_list(object_list, device="cuda")
     return object_list[0]
 
 
+def _get_rank():
+    if dist.is_initialized():
+        return dist.get_rank()
+    else:
+        return 0
+
+
 def get_truncated_value(value):
     if value is None:
         return None

sglang/srt/debug_utils/text_comparator.py

@@ -1,4 +1,5 @@
 import argparse
+import hashlib
 import json
 from pathlib import Path
 
@@ -13,7 +14,11 @@ Supported inputs:
 
 
 def main(args):
-    df_input = _transform_df_input(_compute_df_raw(args))
+    if args.data_type == "simple_evals":
+        df_input = _compute_df_input_mode_simple_evals(args)
+    else:
+        df_input = _transform_df_input(_compute_df_raw(args))
+
     assert all(
         c in df_input.columns
         for c in ["category", "trial_index", "prompt_id", "prompt", "output", "correct"]
@@ -37,8 +42,9 @@ def main(args):
                 df_meta=df_meta.to_dicts(),
                 df_good_to_bad=df_good_to_bad.to_dicts(),
                 df_bad_to_good=df_bad_to_good.to_dicts(),
-            )
-        )
+            ),
+            indent=4,
+        ),
     )
 
     if not args.disable_print_details:
@@ -65,19 +71,70 @@ def main(args):
                 print(df)
 
 
+def _compute_df_input_mode_simple_evals(args):
+    return pl.concat(
+        [
+            _compute_df_input_one_mode_simple_evals(**info)
+            for info in _get_file_infos(args=args)
+        ]
+    )
+
+
+def _compute_df_input_one_mode_simple_evals(path, category, trial_index):
+    data = json.loads(Path(path).read_text())
+    rows = []
+
+    for single_eval_result in data["metadata"]["single_eval_results"]:
+        prompt = single_eval_result["example_level_metadata"][
+            "actual_queried_prompt_messages"
+        ]
+        score = single_eval_result["score"]
+        assert score in {0.0, 1.0}, f"{score=}"
+
+        row = dict(
+            category=category,
+            trial_index=trial_index,
+            prompt_id=_compute_id_from_object(prompt),
+            prompt=json.dumps(prompt),
+            output=single_eval_result["example_level_metadata"]["response_text"],
+            correct=score == 1.0,
+        )
+        rows.append(row)
+
+    return pl.DataFrame(rows)
+
+
+def _compute_id_from_object(obj):
+    if isinstance(obj, pl.Series):
+        obj = obj.to_list()
+    json_str = json.dumps(obj, sort_keys=True, ensure_ascii=False)
+    return hashlib.sha256(json_str.encode("utf-8")).hexdigest()
+
+
 def _compute_df_raw(args):
     return pl.concat(
         [
-            _read_df_raw(p, category=category, trial_index=i)
-            for category, paths in [
-                ("baseline", args.baseline_path),
-                ("target", args.target_path),
-            ]
-            for i, p in enumerate(paths)
+            _read_df_raw(
+                path=info["path"],
+                category=info["category"],
+                trial_index=info["trial_index"],
+            )
+            for info in _get_file_infos(args=args)
         ]
     )
 
 
+def _get_file_infos(args):
+    return [
+        dict(path=path, category=category, trial_index=trial_index)
+        for category, paths in [
+            ("baseline", args.baseline_path),
+            ("target", args.target_path),
+        ]
+        for trial_index, path in enumerate(paths)
+    ]
+
+
 def _read_df_raw(path: str, category: str, trial_index: int):
     return pl.read_ndjson(path).with_columns(
         category=pl.lit(category), trial_index=trial_index
@@ -108,7 +165,9 @@ def _transform_df_input(df: pl.DataFrame):
         print("Transform mode: SGLang bench")
         return df
     else:
-        raise Exception(f"Unknown data: {df.columns}")
+        raise Exception(
+            f"Unknown data: {df.columns}. You may need to set `--data-type` if using e.g. simple_evals."
+        )
 
 
 def _compute_df_meta(df_input: pl.DataFrame):
@@ -127,7 +186,9 @@ def _compute_df_meta(df_input: pl.DataFrame):
 
 
 def _handle_one_prompt(df_one_prompt: pl.DataFrame):
-    assert len(set(df_one_prompt["prompt"])) == 1
+    assert (
+        len(set(_compute_id_from_object(obj) for obj in df_one_prompt["prompt"])) == 1
+    )
 
     df_baseline = df_one_prompt.filter(pl.col("category") == "baseline")
     df_target = df_one_prompt.filter(pl.col("category") == "target")
@@ -162,6 +223,7 @@ def _compute_str_prefix_len(a: str, b: str) -> int:
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description=_DESCRIPTION)
+    parser.add_argument("--data-type", type=str, default="auto")
     parser.add_argument("--baseline-path", type=str, nargs="+")
     parser.add_argument("--target-path", type=str, nargs="+")
     parser.add_argument(

sglang/srt/disaggregation/ascend/conn.py

@@ -13,7 +13,7 @@ from sglang.srt.disaggregation.mooncake.conn import (
     MooncakeKVReceiver,
     MooncakeKVSender,
 )
-from sglang.srt.utils import get_local_ip_by_remote
+from sglang.srt.utils import get_local_ip_auto
 
 logger = logging.getLogger(__name__)
 
@@ -21,7 +21,7 @@ logger = logging.getLogger(__name__)
 class AscendKVManager(MooncakeKVManager):
     def init_engine(self):
         # TransferEngine initialized on ascend.
-        local_ip = get_local_ip_by_remote()
+        local_ip = get_local_ip_auto()
         self.engine = AscendTransferEngine(
             hostname=local_ip,
             npu_id=self.kv_args.gpu_id,

sglang/srt/disaggregation/ascend/transfer_engine.py

@@ -2,9 +2,19 @@ import logging
 import os
 from typing import List, Optional
 
+import torch
+
 from sglang.srt.disaggregation.mooncake.transfer_engine import MooncakeTransferEngine
 from sglang.srt.disaggregation.utils import DisaggregationMode
 
+try:
+    from mf_adapter import TransferEngine
+
+    import_error = None
+except ImportError as e:
+    import_error = e
+    pass
+
 logger = logging.getLogger(__name__)
 
 
@@ -13,12 +23,11 @@ class AscendTransferEngine(MooncakeTransferEngine):
     def __init__(
         self, hostname: str, npu_id: int, disaggregation_mode: DisaggregationMode
     ):
-        try:
-            from mf_adapter import TransferEngine
-        except ImportError as e:
-            raise ImportError(
+        if import_error is not None:
+            logger.warning(
                 "Please install mf_adapter, for details, see docs/backend/pd_disaggregation.md"
-            ) from e
+            )
+            raise import_error
 
         self.engine = TransferEngine()
         self.hostname = hostname
@@ -37,12 +46,29 @@ class AscendTransferEngine(MooncakeTransferEngine):
         self.initialize()
 
     def initialize(self) -> None:
+        from sglang.srt.layers.dp_attention import (
+            get_tensor_model_parallel_world_size,
+            get_tp_group,
+        )
+
+        transfer_protocol = self._get_transfer_protocol()
+        if transfer_protocol is None or transfer_protocol == "sdma":
+            trans_op_type = TransferEngine.TransDataOpType.SDMA
+        else:
+            trans_op_type = TransferEngine.TransDataOpType.DEVICE_RDMA
+            """with device RDMA for PD transfer"""
+            tmp_tensor = torch.zeros(1, device="npu")
+            output_tensor_list = [
+                torch.empty_like(tmp_tensor)
+                for _ in range(get_tensor_model_parallel_world_size())
+            ]
+            # Initialize hccl in advance through all_gather to avoid conflicts with rdma initialization.
+            torch.distributed.all_gather(
+                output_tensor_list, tmp_tensor, group=get_tp_group().device_group
+            )
         """Initialize the ascend transfer instance."""
         ret_value = self.engine.initialize(
-            self.store_url,
-            self.session_id,
-            self.role,
-            self.npu_id,
+            self.store_url, self.session_id, self.role, self.npu_id, trans_op_type
         )
         if ret_value != 0:
             logger.error("Ascend Transfer Engine initialization failed.")
@@ -56,3 +82,15 @@ class AscendTransferEngine(MooncakeTransferEngine):
             ret_value = -1
         if ret_value != 0:
             logger.debug(f"Ascend memory registration for ptr {ptrs} failed.")
+
+    @staticmethod
+    def _get_transfer_protocol():
+        protocol = os.getenv("ASCEND_MF_TRANSFER_PROTOCOL")
+        allowed_protocols = {"device_rdma", "sdma"}
+        if protocol and protocol.lower() in allowed_protocols:
+            return protocol.lower()
+        else:
+            logger.warning(
+                "Invalid or no transfer protocol specified, using default protocol."
+            )
+            return None

sglang/srt/disaggregation/base/conn.py

@@ -131,4 +131,4 @@ class BaseKVReceiver(ABC):
 
 class BaseKVBootstrapServer(ABC):
     @abstractmethod
-    def __init__(self, port: int): ...
+    def __init__(self, host: str, port: int): ...