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

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/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): ...

sglang/srt/disaggregation/common/conn.py

@@ -47,6 +47,7 @@ class CommonKVManager(BaseKVManager):
         self.is_mla_backend = is_mla_backend
         self.disaggregation_mode = disaggregation_mode
         # for p/d multi node infer
+        self.bootstrap_host = server_args.host
         self.bootstrap_port = server_args.disaggregation_bootstrap_port
         self.dist_init_addr = server_args.dist_init_addr
         self.tp_size = server_args.tp_size
@@ -72,6 +73,7 @@ class CommonKVManager(BaseKVManager):
     def _register_to_bootstrap(self):
         """Register KVSender to bootstrap server via HTTP POST."""
         if self.dist_init_addr:
+            # multi node: bootstrap server's host is dist_init_addr
             if self.dist_init_addr.startswith("["):  # [ipv6]:port or [ipv6]
                 if self.dist_init_addr.endswith("]"):
                     host = self.dist_init_addr
@@ -80,7 +82,8 @@ class CommonKVManager(BaseKVManager):
             else:
                 host = socket.gethostbyname(self.dist_init_addr.rsplit(":", 1)[0])
         else:
-            host = get_ip()
+            # single node: bootstrap server's host is same as http server's host
+            host = self.bootstrap_host
             host = maybe_wrap_ipv6_address(host)
 
         bootstrap_server_url = f"{host}:{self.bootstrap_port}"
@@ -125,12 +128,11 @@ class CommonKVReceiver(BaseKVReceiver):
         mgr: BaseKVManager,
         bootstrap_addr: str,
         bootstrap_room: Optional[int] = None,
-        data_parallel_rank: Optional[int] = None,
+        prefill_dp_rank: Optional[int] = None,
     ):
         self.bootstrap_room = bootstrap_room
         self.bootstrap_addr = bootstrap_addr
         self.kv_mgr = mgr
-        self.data_parallel_rank = data_parallel_rank
 
         if self.bootstrap_addr not in self.kv_mgr.prefill_dp_size_table:
             self.prefill_tp_size, self.prefill_dp_size = (
@@ -166,9 +168,6 @@ class CommonKVReceiver(BaseKVReceiver):
             self.required_dst_info_num = 1
             self.target_tp_ranks = [self.target_tp_rank]
         elif local_tp_size_per_dp_rank > prefill_tp_size_per_dp_rank:
-            assert (
-                self.kv_mgr.is_mla_backend
-            ), "PD with different TP sizes per DP rank is not yet supported for non-MLA models"
             self.target_tp_rank = (
                 self.kv_mgr.kv_args.engine_rank % local_tp_size_per_dp_rank
             ) // (local_tp_size_per_dp_rank // prefill_tp_size_per_dp_rank)
@@ -198,11 +197,14 @@ class CommonKVReceiver(BaseKVReceiver):
             self.target_tp_rank = self.target_tp_ranks[0]
             self.required_dst_info_num = 1
 
-        if self.data_parallel_rank is not None:
-            logger.debug(f"Targeting DP rank: {self.data_parallel_rank}")
-            self.target_dp_group = self.data_parallel_rank
+        if prefill_dp_rank is not None:
+            logger.debug(f"Targeting DP rank: {prefill_dp_rank}")
+            self.prefill_dp_rank = prefill_dp_rank
         else:
-            self.target_dp_group = bootstrap_room % self.prefill_dp_size
+            self.prefill_dp_rank = bootstrap_room % self.prefill_dp_size
+
+        # FIXME: alias here: target_dp_group -> prefill_dp_rank
+        self.target_dp_group = self.prefill_dp_rank
 
         # NOTE: key distinguished by bootstrap_addr, target_dp_group, and target_tp_rank
         bootstrap_key = (
@@ -308,7 +310,8 @@ class CommonKVReceiver(BaseKVReceiver):
 
 
 class CommonKVBootstrapServer(BaseKVBootstrapServer):
-    def __init__(self, port: int):
+    def __init__(self, host: str, port: int):
+        self.host = host
         self.port = port
         self.app = web.Application()
         self.store = dict()
@@ -412,7 +415,7 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
             self._runner = web.AppRunner(self.app)
             self._loop.run_until_complete(self._runner.setup())
 
-            site = web.TCPSite(self._runner, port=self.port)
+            site = web.TCPSite(self._runner, host=self.host, port=self.port)
             self._loop.run_until_complete(site.start())
             self._loop.run_forever()
         except Exception as e:

sglang/srt/disaggregation/decode.py

@@ -24,7 +24,7 @@ import logging
 from collections import deque
 from dataclasses import dataclass
 from http import HTTPStatus
-from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union
 
 import torch
 from torch.distributed import ProcessGroup
@@ -218,8 +218,10 @@ class DecodePreallocQueue:
 
         kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device
         kv_args.gpu_id = self.scheduler.gpu_id
-        kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER)
-        kv_manager = kv_manager_class(
+        kv_manager_class: Type[BaseKVManager] = get_kv_class(
+            self.transfer_backend, KVClassType.MANAGER
+        )
+        kv_manager: BaseKVManager = kv_manager_class(
             kv_args,
             DisaggregationMode.DECODE,
             self.scheduler.server_args,
@@ -248,7 +250,7 @@ class DecodePreallocQueue:
                 mgr=self.kv_manager,
                 bootstrap_addr=f"{req.bootstrap_host}:{req.bootstrap_port}",
                 bootstrap_room=req.bootstrap_room,
-                data_parallel_rank=req.data_parallel_rank,
+                prefill_dp_rank=req.data_parallel_rank,
             )
 
             self.queue.append(
@@ -884,9 +886,18 @@ class SchedulerDisaggregationDecodeMixin:
             # if there are still retracted requests, we do not allocate new requests
             return
 
-        req_conns = self.disagg_decode_prealloc_queue.pop_preallocated()
-        self.disagg_decode_transfer_queue.extend(req_conns)
-        alloc_reqs = (
-            self.disagg_decode_transfer_queue.pop_transferred()
-        )  # the requests which kv has arrived
-        self.waiting_queue.extend(alloc_reqs)
+        if not hasattr(self, "polling_count"):
+            self.polling_count = 0
+            self.polling_interval = (
+                self.server_args.disaggregation_decode_polling_interval
+            )
+
+        self.polling_count = (self.polling_count + 1) % self.polling_interval
+
+        if self.polling_count % self.polling_interval == 0:
+            req_conns = self.disagg_decode_prealloc_queue.pop_preallocated()
+            self.disagg_decode_transfer_queue.extend(req_conns)
+            alloc_reqs = (
+                self.disagg_decode_transfer_queue.pop_transferred()
+            )  # the requests which kv has arrived
+            self.waiting_queue.extend(alloc_reqs)

sglang/srt/disaggregation/decode_schedule_batch_mixin.py

@@ -110,7 +110,10 @@ class ScheduleBatchDisaggregationDecodeMixin:
             if req.grammar is not None:
                 # FIXME: this try-except block is for handling unexpected xgrammar issue.
                 try:
-                    req.grammar.accept_token(req.output_ids[-1])
+                    # if it is not None, then the grammar is from a retracted request, and we should not
+                    # accept the token as it's already accepted
+                    if req.grammar.current_token is None:
+                        req.grammar.accept_token(req.output_ids[-1])
                 except ValueError as e:
                     # Grammar accept_token can raise ValueError if the token is not in the grammar.
                     # This can happen if the grammar is not set correctly or the token is invalid.

sglang/srt/disaggregation/fake/conn.py

@@ -62,7 +62,7 @@ class FakeKVReceiver(BaseKVReceiver):
         mgr: BaseKVManager,
         bootstrap_addr: str,
         bootstrap_room: Optional[int] = None,
-        data_parallel_rank: Optional[int] = None,
+        prefill_dp_rank: Optional[int] = None,
     ):
         self.has_init = False