sglang 0.5.1.post2__py3-none-any.whl → 0.5.2__py3-none-any.whl

sglang/srt/debug_utils/dump_comparator.py

@@ -1,11 +1,11 @@
 import argparse
 import functools
-import re
 from pathlib import Path
 
 import polars as pl
 import torch
 
+from sglang.srt.debug_utils.dump_loader import find_row, read_meta
 from sglang.srt.debug_utils.dumper import get_truncated_value
 
 
@@ -26,66 +26,77 @@ def main(args):
     print("df_baseline", df_baseline)
 
     for row in df_target.iter_rows(named=True):
-        rows_baseline = df_baseline.filter(
-            (
-                pl.col("forward_pass_id")
-                == row["forward_pass_id"] - args.start_id + args.baseline_start_id
-            )
-            & functools.reduce(
-                lambda a, b: a & b,
-                [
-                    pl.col(col) == row[col]
-                    for col in row.keys()
-                    if col not in ["forward_pass_id", "dump_index", "filename"]
-                ],
-            )
+        path_target = Path(args.target_path) / row["filename"]
+
+        row_baseline = find_row(
+            df_baseline,
+            conditions=dict(
+                forward_pass_id=row["forward_pass_id"]
+                - args.start_id
+                + args.baseline_start_id,
+                **{
+                    k: v
+                    for k, v in row.items()
+                    if k not in ["forward_pass_id", "dump_index", "filename"]
+                },
+            ),
         )
-        assert len(rows_baseline) == 1, f"{rows_baseline=}"
-        row_baseline = rows_baseline.to_dicts()[0]
+
+        if row_baseline is None:
+            print(f"Skip: target={str(path_target)} since no baseline")
+            x_target = _load_object(path_target)
+            if x_target is not None:
+                print(f"x_target(sample)={get_truncated_value(x_target)}")
+            continue
 
         path_baseline = Path(args.baseline_path) / row_baseline["filename"]
-        path_target = Path(args.target_path) / row["filename"]
         print(f"Check: target={str(path_target)} baseline={str(path_baseline)}")
-        check_tensor_pair(path_baseline=path_baseline, path_target=path_target)
+        check_tensor_pair(
+            path_baseline=path_baseline, path_target=path_target, name=row["name"]
+        )
         print()
 
 
-def read_meta(directory):
-    directory = Path(directory)
-    assert directory.is_dir(), f"{directory=} should be a directory"
-
-    rows = []
-    for p in directory.glob("*.pt"):
-        full_kwargs = {}
-        for kv in p.stem.split("___"):
-            k, v = kv.split("=")
-            full_kwargs[k] = v
-        rows.append(
-            {
-                "filename": str(p.name),
-                **full_kwargs,
-            }
-        )
+def check_tensor_pair(path_baseline, path_target, name=""):
+    x_baseline = _load_object(path_baseline)
+    x_target = _load_object(path_target)
 
-    df = pl.DataFrame(rows)
-    df = df.with_columns(
-        pl.col("forward_pass_id").cast(int),
-        pl.col("rank").cast(int),
+    print(
+        f"Raw "
+        f"[shape] {x_baseline.shape} vs {x_target.shape}\t"
+        f"[dtype] {x_baseline.dtype} vs {x_target.dtype}"
     )
-    return df
-
 
-def check_tensor_pair(path_baseline, path_target):
-    x_baseline = torch.load(path_baseline, weights_only=True)
-    x_target = torch.load(path_target, weights_only=True)
+    x_baseline, x_target = _comparison_preprocessor(x_baseline, x_target, name=name)
+    x_baseline = _try_unify_shape(x_baseline, target_shape=x_target.shape)
 
     print(
+        f"After preprocessor "
         f"[shape] {x_baseline.shape} vs {x_target.shape}\t"
         f"[dtype] {x_baseline.dtype} vs {x_target.dtype}"
     )
 
+    x_target = x_target.float()
+    x_baseline = x_baseline.float()
+
+    for name, fn in (
+        ("mean", torch.mean),
+        ("std", torch.std),
+        ("min", torch.min),
+        ("max", torch.max),
+        ("p1", functools.partial(torch.quantile, q=0.01)),
+        ("p5", functools.partial(torch.quantile, q=0.05)),
+        ("p95", functools.partial(torch.quantile, q=0.95)),
+        ("p99", functools.partial(torch.quantile, q=0.99)),
+    ):
+        value_baseline = fn(x_baseline).item()
+        value_target = fn(x_target).item()
+        print(
+            f"[{name}] {value_baseline :.4f} vs {value_target:.4f} (diff: {value_target - value_baseline:.4f})"
+        )
+
     if x_baseline.shape != x_target.shape:
-        print(f"❌ Shape mismatch")
+        print(f"⚠️ Shape mismatch")
         return
 
     raw_abs_diff = (x_target - x_baseline).abs()
@@ -112,6 +123,19 @@ def check_tensor_pair(path_baseline, path_target):
         print(f"x_target(sample)={get_truncated_value(x_target)}")
 
 
+def _try_unify_shape(x: torch.Tensor, target_shape):
+    x_shape = x.shape
+    num_dim_to_remove = len(x_shape) - len(target_shape)
+    if (x_shape[num_dim_to_remove:] == target_shape) and all(
+        val == 1 for val in x_shape[:num_dim_to_remove]
+    ):
+        out = functools.reduce(lambda a, _: a.squeeze(0), range(num_dim_to_remove), x)
+        print(f"Unify shape: {x_shape} -> {out.shape} (to match {target_shape})")
+        return out
+
+    return x
+
+
 # Copied from DeepGEMM
 def _calc_rel_diff(x: torch.Tensor, y: torch.Tensor):
     x, y = x.double(), y.double()
@@ -120,6 +144,19 @@ def _calc_rel_diff(x: torch.Tensor, y: torch.Tensor):
     return 1 - sim
 
 
+def _comparison_preprocessor(x_baseline, x_target, name):
+    # can insert arbitrary adhoc postprocessing logic here
+    return x_baseline, x_target
+
+
+def _load_object(path):
+    x = torch.load(path, weights_only=False)
+    if not isinstance(x, torch.Tensor):
+        print(f"Skip load {path} since {type(x)=} is not a Tensor")
+        return None
+    return x.cuda()
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--baseline-path", type=str)

sglang/srt/debug_utils/dump_loader.py

@@ -0,0 +1,97 @@
+import functools
+import os
+from pathlib import Path
+from typing import Any, Dict
+
+import polars as pl
+import torch
+
+
+class DumpLoader:
+    def __init__(self):
+        directory = os.environ.get("SGLANG_DUMP_LOADER_DIR")
+
+        self._enable = directory is not None
+        if self._enable:
+            self._directory = Path(directory)
+            self._df = read_meta(directory)
+
+    @property
+    def enable(self):
+        return self._enable
+
+    def load(self, name, **kwargs):
+        assert self._enable, "Please call DumpLoader.load only when it is enabled"
+
+        from sglang.srt.debug_utils.dumper import dumper
+
+        forward_pass_id = dumper._forward_pass_id
+        conditions = dict(name=name, forward_pass_id=forward_pass_id, **kwargs)
+        row = find_row(self._df, conditions=conditions)
+        assert (
+            row is not None
+        ), f"DumpLoader cannot find row given query {name=} {kwargs=} {self._directory=}"
+
+        path = self._directory / row["filename"]
+        output = torch.load(path, weights_only=False)
+
+        print(
+            f"[DumpLoader] load from {path=} (query: {name=} {kwargs=}, output: {type(output)})"
+        )
+        return output
+
+
+def read_meta(directory):
+    directory = Path(directory)
+    assert directory.is_dir(), f"{directory=} should be a directory"
+
+    rows = []
+    for p in directory.glob("*.pt"):
+        full_kwargs = {}
+        for kv in p.stem.split("___"):
+            k, v = kv.split("=")
+            full_kwargs[k] = v
+        rows.append(
+            {
+                "filename": str(p.name),
+                **full_kwargs,
+            }
+        )
+
+    df = pl.DataFrame(rows)
+    df = df.with_columns(
+        pl.col("forward_pass_id").cast(int),
+        pl.col("rank").cast(int),
+        pl.col("dump_index").cast(int),
+    )
+    return df
+
+
+def find_row(df, conditions: Dict[str, Any]):
+    df_sub = df.filter(
+        functools.reduce(
+            lambda a, b: a & b,
+            [
+                pl.col(col) == _cast_to_polars_dtype(conditions[col], df.schema[col])
+                for col in conditions.keys()
+            ],
+        )
+    )
+    assert len(df_sub) <= 1
+    return df_sub.to_dicts()[0] if len(df_sub) > 0 else None
+
+
+def _cast_to_polars_dtype(value, target_dtype):
+    if target_dtype in (pl.Int64, pl.Int32, pl.UInt64, pl.UInt32):
+        return int(value)
+    elif target_dtype in (pl.Float64, pl.Float32):
+        return float(value)
+    elif target_dtype == pl.Boolean:
+        return bool(value)
+    elif target_dtype == pl.String:
+        return str(value)
+    else:
+        return value
+
+
+dump_loader = DumpLoader()

sglang/srt/debug_utils/dumper.py

@@ -53,7 +53,7 @@ class _Dumper:
         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 +80,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

@@ -1,6 +1,12 @@
+import concurrent.futures
 import logging
+from typing import List, Tuple
+
+import numpy as np
+import numpy.typing as npt
 
 from sglang.srt.disaggregation.ascend.transfer_engine import AscendTransferEngine
+from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous
 from sglang.srt.disaggregation.mooncake.conn import (
     MooncakeKVBootstrapServer,
     MooncakeKVManager,
@@ -29,6 +35,75 @@ class AscendKVManager(MooncakeKVManager):
             self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens
         )
 
+    def send_kvcache(
+        self,
+        mooncake_session_id: str,
+        prefill_kv_indices: npt.NDArray[np.int32],
+        dst_kv_ptrs: list[int],
+        dst_kv_indices: npt.NDArray[np.int32],
+        executor: concurrent.futures.ThreadPoolExecutor,
+    ):
+        # Group by indices
+        prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
+            prefill_kv_indices, dst_kv_indices
+        )
+
+        num_layers = len(self.kv_args.kv_data_ptrs)
+        layers_params = [
+            (
+                self.kv_args.kv_data_ptrs[layer_id],
+                dst_kv_ptrs[layer_id],
+                self.kv_args.kv_item_lens[layer_id],
+            )
+            for layer_id in range(num_layers)
+        ]
+
+        def set_transfer_blocks(
+            src_ptr: int, dst_ptr: int, item_len: int
+        ) -> List[Tuple[int, int, int]]:
+            transfer_blocks = []
+            for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks):
+                src_addr = src_ptr + int(prefill_index[0]) * item_len
+                dst_addr = dst_ptr + int(decode_index[0]) * item_len
+                length = item_len * len(prefill_index)
+                transfer_blocks.append((src_addr, dst_addr, length))
+            return transfer_blocks
+
+        # Worker function for processing a single layer
+        def process_layer(src_ptr: int, dst_ptr: int, item_len: int) -> int:
+            transfer_blocks = set_transfer_blocks(src_ptr, dst_ptr, item_len)
+            return self._transfer_data(mooncake_session_id, transfer_blocks)
+
+        # Worker function for processing all layers in a batch
+        def process_layers(layers_params: List[Tuple[int, int, int]]) -> int:
+            transfer_blocks = []
+            for src_ptr, dst_ptr, item_len in layers_params:
+                transfer_blocks.extend(set_transfer_blocks(src_ptr, dst_ptr, item_len))
+            return self._transfer_data(mooncake_session_id, transfer_blocks)
+
+        if self.enable_custom_mem_pool:
+            futures = [
+                executor.submit(
+                    process_layer,
+                    src_ptr,
+                    dst_ptr,
+                    item_len,
+                )
+                for (src_ptr, dst_ptr, item_len) in layers_params
+            ]
+            for future in concurrent.futures.as_completed(futures):
+                status = future.result()
+                if status != 0:
+                    for f in futures:
+                        f.cancel()
+                    return status
+        else:
+            # Combining all layers' params in one batch transfer is more efficient
+            # compared to using multiple threads
+            return process_layers(layers_params)
+
+        return 0
+
 
 class AscendKVSender(MooncakeKVSender):
     pass

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(

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