sglang 0.4.6.post5__py3-none-any.whl → 0.4.7__py3-none-any.whl

sglang/srt/disaggregation/kv_events.py

@@ -43,6 +43,7 @@ class EventBatch(
 ):
     ts: float
     events: list[Any]
+    attn_dp_rank: Optional[int] = None
 
 
 class KVCacheEvent(
@@ -76,7 +77,21 @@ class KVEventBatch(EventBatch):
 
 
 class EventPublisher(ABC):
-    """Lightweight publisher for EventBatch batches."""
+    """
+    Lightweight publisher for EventBatch batches with
+    support for DP attention.
+
+    In DP attention - each rank has its own Scheduler and
+    KV cache instance in order to avoid duplicate events
+    and ensure proper event attribution. In our implementation
+
+    - Each DP rank has its own EventPublisher
+    - Publishers annotate events with the dp rank
+    - This allows consumers to distinguish events from different DP ranks
+    """
+
+    def __init__(self, attn_dp_rank: int = 0):
+        self._attn_dp_rank = attn_dp_rank
 
     @abstractmethod
     def publish(self, events: EventBatch) -> None:
@@ -130,6 +145,7 @@ class ZmqEventPublisher(EventPublisher):
 
     def __init__(
         self,
+        attn_dp_rank: int,
         endpoint: str = "tcp://*:5557",
         replay_endpoint: Optional[str] = None,
         buffer_steps: int = 10_000,
@@ -138,6 +154,7 @@ class ZmqEventPublisher(EventPublisher):
         topic: str = "",
     ) -> None:
         # Storage
+        super().__init__(attn_dp_rank)
         self._event_queue = Queue[Optional[EventBatch]](maxsize=max_queue_size)
         self._buffer = deque[tuple[int, bytes]](maxlen=buffer_steps)
 
@@ -145,8 +162,11 @@ class ZmqEventPublisher(EventPublisher):
         self._ctx = zmq.Context.instance()
         self._pub: Optional[zmq.Socket] = None
         self._replay: Optional[zmq.Socket] = None
-        self._endpoint = endpoint
-        self._replay_endpoint = replay_endpoint
+        self._dp_rank = attn_dp_rank
+        self._endpoint = self.offset_endpoint_port(endpoint, self._dp_rank)
+        self._replay_endpoint = self.offset_endpoint_port(
+            replay_endpoint, self._dp_rank
+        )
         self._hwm = hwm
         self._socket_setup()
 
@@ -168,6 +188,8 @@ class ZmqEventPublisher(EventPublisher):
     def publish(self, events: EventBatch) -> None:
         if not self._running:
             raise RuntimeError("Publisher is closed")
+        if events.attn_dp_rank is None:
+            events.attn_dp_rank = self._dp_rank
         self._event_queue.put(events)
 
     def shutdown(self) -> None:
@@ -288,6 +310,39 @@ class ZmqEventPublisher(EventPublisher):
         # receiving payload is (-1, b""")
         self._replay.send_multipart((client_id, b"", self.END_SEQ, b""))
 
+    @staticmethod
+    def offset_endpoint_port(
+        endpoint: Optional[str], data_parallel_rank: int
+    ) -> Optional[str]:
+        """Helper function to offset the port in an endpoint by
+            the data parallel rank.
+
+        Args:
+            endpoint: The endpoint string
+                (e.g., "tcp://*:5557" or "inproc://cache")
+            data_parallel_rank: The data parallel rank to offset by
+
+        Returns:
+            The endpoint with the port offset by data_parallel_rank
+                or suffix appended
+        """
+        # Do nothing if input is None or data_parallel_rank is 0
+        if not endpoint or data_parallel_rank == 0:
+            return endpoint
+
+        if "inproc" in endpoint:
+            return f"{endpoint}_dp{data_parallel_rank}"
+        if "tcp" in endpoint:
+            if endpoint and ":" in endpoint:
+                # Get everything after the last colon (the port)
+                last_colon_idx = endpoint.rfind(":")
+                base_addr = endpoint[:last_colon_idx]
+                base_port = int(endpoint[last_colon_idx + 1 :])
+                new_port = base_port + data_parallel_rank
+                return f"{base_addr}:{new_port}"
+            return endpoint
+        raise ValueError("Invalid endpoint: must contain 'inproc' or 'tcp'")
+
 
 class KVEventsConfig(BaseModel):
     """Configuration for KV event publishing."""
@@ -342,7 +397,7 @@ class EventPublisherFactory:
         cls._registry[name] = ctor
 
     @classmethod
-    def create(cls, config: Optional[str]) -> EventPublisher:
+    def create(cls, config: Optional[str], attn_dp_rank: int = 0) -> EventPublisher:
         """Create publisher from a config mapping."""
         if not config:
             return NullEventPublisher()
@@ -354,4 +409,4 @@ class EventPublisherFactory:
             constructor = cls._registry[kind]
         except KeyError as exc:
             raise ValueError(f"Unknown event publisher '{kind}'") from exc
-        return constructor(**config_dict)
+        return constructor(attn_dp_rank=attn_dp_rank, **config_dict)

sglang/srt/disaggregation/launch_lb.py

@@ -0,0 +1,140 @@
+import argparse
+import dataclasses
+
+
+@dataclasses.dataclass
+class LBArgs:
+    rust_lb: bool = False
+    host: str = "0.0.0.0"
+    port: int = 8000
+    policy: str = "random"
+    prefill_infos: list = dataclasses.field(default_factory=list)
+    decode_infos: list = dataclasses.field(default_factory=list)
+    log_interval: int = 5
+    timeout: int = 600
+
+    @staticmethod
+    def add_cli_args(parser: argparse.ArgumentParser):
+        parser.add_argument(
+            "--rust-lb",
+            action="store_true",
+            help="Use Rust load balancer",
+        )
+        parser.add_argument(
+            "--host",
+            type=str,
+            default=LBArgs.host,
+            help=f"Host to bind the server (default: {LBArgs.host})",
+        )
+        parser.add_argument(
+            "--port",
+            type=int,
+            default=LBArgs.port,
+            help=f"Port to bind the server (default: {LBArgs.port})",
+        )
+        parser.add_argument(
+            "--policy",
+            type=str,
+            default=LBArgs.policy,
+            choices=["random", "po2"],
+            help=f"Policy to use for load balancing (default: {LBArgs.policy})",
+        )
+        parser.add_argument(
+            "--prefill",
+            type=str,
+            default=[],
+            nargs="+",
+            help="URLs for prefill servers",
+        )
+        parser.add_argument(
+            "--decode",
+            type=str,
+            default=[],
+            nargs="+",
+            help="URLs for decode servers",
+        )
+        parser.add_argument(
+            "--prefill-bootstrap-ports",
+            type=int,
+            nargs="+",
+            help="Bootstrap ports for prefill servers",
+        )
+        parser.add_argument(
+            "--log-interval",
+            type=int,
+            default=LBArgs.log_interval,
+            help=f"Log interval in seconds (default: {LBArgs.log_interval})",
+        )
+        parser.add_argument(
+            "--timeout",
+            type=int,
+            default=LBArgs.timeout,
+            help=f"Timeout in seconds (default: {LBArgs.timeout})",
+        )
+
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace) -> "LBArgs":
+        bootstrap_ports = args.prefill_bootstrap_ports
+        if bootstrap_ports is None:
+            bootstrap_ports = [None] * len(args.prefill)
+        elif len(bootstrap_ports) == 1:
+            bootstrap_ports = bootstrap_ports * len(args.prefill)
+        else:
+            if len(bootstrap_ports) != len(args.prefill):
+                raise ValueError(
+                    "Number of prefill URLs must match number of bootstrap ports"
+                )
+
+        prefill_infos = [
+            (url, port) for url, port in zip(args.prefill, bootstrap_ports)
+        ]
+
+        return cls(
+            rust_lb=args.rust_lb,
+            host=args.host,
+            port=args.port,
+            policy=args.policy,
+            prefill_infos=prefill_infos,
+            decode_infos=args.decode,
+            log_interval=args.log_interval,
+            timeout=args.timeout,
+        )
+
+    def __post_init__(self):
+        if not self.rust_lb:
+            assert (
+                self.policy == "random"
+            ), "Only random policy is supported for Python load balancer"
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="PD Disaggregation Load Balancer Server"
+    )
+    LBArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    lb_args = LBArgs.from_cli_args(args)
+
+    if lb_args.rust_lb:
+        from sgl_pdlb._rust import LoadBalancer as RustLB
+
+        RustLB(
+            host=lb_args.host,
+            port=lb_args.port,
+            policy=lb_args.policy,
+            prefill_infos=lb_args.prefill_infos,
+            decode_infos=lb_args.decode_infos,
+            log_interval=lb_args.log_interval,
+            timeout=lb_args.timeout,
+        ).start()
+    else:
+        from sglang.srt.disaggregation.mini_lb import PrefillConfig, run
+
+        prefill_configs = [
+            PrefillConfig(url, port) for url, port in lb_args.prefill_infos
+        ]
+        run(prefill_configs, lb_args.decode_infos, lb_args.host, lb_args.port)
+
+
+if __name__ == "__main__":
+    main()

sglang/srt/disaggregation/mini_lb.py

@@ -50,6 +50,13 @@ class MiniLoadBalancer:
         self.prefill_servers = [p.url for p in prefill_configs]
         self.decode_servers = decode_servers
 
+    def add_prefill_server(self, new_prefill_config: PrefillConfig):
+        self.prefill_configs.append(new_prefill_config)
+        self.prefill_servers.append(new_prefill_config.url)
+
+    def add_decode_server(self, new_decode_server: str):
+        self.decode_servers.append(new_decode_server)
+
     def select_pair(self):
         # TODO: return some message instead of panic
         assert len(self.prefill_configs) > 0, "No prefill servers available"
@@ -110,8 +117,8 @@ class MiniLoadBalancer:
             ) as session:
                 # Create the tasks for both prefill and decode requests
                 tasks = [
-                    session.post(f"{prefill_server}/generate", json=modified_request),
-                    session.post(f"{decode_server}/generate", json=modified_request),
+                    session.post(f"{prefill_server}/{endpoint}", json=modified_request),
+                    session.post(f"{decode_server}/{endpoint}", json=modified_request),
                 ]
                 # Wait for both responses to complete. Since this is streaming, they return immediately.
                 prefill_response, decode_response = await asyncio.gather(*tasks)
@@ -157,7 +164,7 @@ class MiniLoadBalancer:
 
 
 app = FastAPI()
-load_balancer = None
+load_balancer: Optional[MiniLoadBalancer] = None
 
 
 @app.get("/health")
@@ -267,8 +274,7 @@ async def handle_generate_request(request_data: dict):
         )
 
 
-@app.post("/v1/chat/completions")
-async def handle_completion_request(request_data: dict):
+async def _forward_to_backend(request_data: dict, endpoint_name: str):
     prefill_server, bootstrap_port, decode_server = load_balancer.select_pair()
 
     # Parse and transform prefill_server for bootstrap data
@@ -279,7 +285,7 @@ async def handle_completion_request(request_data: dict):
         {
             "bootstrap_host": hostname,
             "bootstrap_port": bootstrap_port,
-            "bootstrap_room": random.randint(0, 2**63 - 1),
+            "bootstrap_room": _generate_bootstrap_room(),
         }
     )
 
@@ -288,17 +294,27 @@ async def handle_completion_request(request_data: dict):
             modified_request,
             prefill_server,
             decode_server,
-            endpoint="v1/chat/completions",
+            endpoint=endpoint_name,
         )
     else:
         return await load_balancer.generate(
             modified_request,
             prefill_server,
             decode_server,
-            endpoint="v1/chat/completions",
+            endpoint=endpoint_name,
         )
 
 
+@app.post("/v1/chat/completions")
+async def handle_chat_completion_request(request_data: dict):
+    return await _forward_to_backend(request_data, "v1/chat/completions")
+
+
+@app.post("/v1/completions")
+async def handle_completion_request(request_data: dict):
+    return await _forward_to_backend(request_data, "v1/completions")
+
+
 def _generate_bootstrap_room():
     return random.randint(0, 2**63 - 1)
 
@@ -331,14 +347,14 @@ async def get_models():
 @app.post("/register")
 async def register(obj: PDRegistryRequest):
     if obj.mode == "prefill":
-        load_balancer.prefill_configs.append(
+        load_balancer.add_prefill_server(
             PrefillConfig(obj.registry_url, obj.bootstrap_port)
         )
         logger.info(
             f"Registered prefill server: {obj.registry_url} with bootstrap port: {obj.bootstrap_port}"
         )
     elif obj.mode == "decode":
-        load_balancer.decode_servers.append(obj.registry_url)
+        load_balancer.add_decode_server(obj.registry_url)
         logger.info(f"Registered decode server: {obj.registry_url}")
     else:
         raise HTTPException(
@@ -361,42 +377,7 @@ def run(prefill_configs, decode_addrs, host, port):
 
 
 if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser(description="Mini Load Balancer Server")
-    parser.add_argument(
-        "--prefill", type=str, default=[], nargs="+", help="URLs for prefill servers"
-    )
-    parser.add_argument(
-        "--decode", type=str, default=[], nargs="+", help="URLs for decode servers"
-    )
-    parser.add_argument(
-        "--prefill-bootstrap-ports",
-        type=int,
-        nargs="+",
-        help="Bootstrap ports for prefill servers",
-    )
-    parser.add_argument(
-        "--host", default="0.0.0.0", help="Host to bind the server (default: 0.0.0.0)"
-    )
-    parser.add_argument(
-        "--port", type=int, default=8000, help="Port to bind the server (default: 8000)"
-    )
-    args = parser.parse_args()
-
-    bootstrap_ports = args.prefill_bootstrap_ports
-    if bootstrap_ports is None:
-        bootstrap_ports = [None] * len(args.prefill)
-    elif len(bootstrap_ports) == 1:
-        bootstrap_ports = bootstrap_ports * len(args.prefill)
-    else:
-        if len(bootstrap_ports) != len(args.prefill):
-            raise ValueError(
-                "Number of prefill URLs must match number of bootstrap ports"
-            )
-
-    prefill_configs = [
-        PrefillConfig(url, port) for url, port in zip(args.prefill, bootstrap_ports)
-    ]
+    # FIXME: remove this, use the unified entry point: sglang.srt.disaggregation.launch_lb
+    from sglang.srt.disaggregation.launch_lb import main
 
-    run(prefill_configs, args.decode, args.host, args.port)
+    main()