sglang 0.4.5.post2__py3-none-any.whl → 0.4.6__py3-none-any.whl

sglang/srt/disaggregation/decode.py

@@ -21,6 +21,7 @@ Life cycle of a request in the decode server
 from __future__ import annotations
 
 import logging
+from collections import deque
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
@@ -136,7 +137,7 @@ class DecodePreallocQueue:
         kv_receiver_class = get_kv_class(self.transfer_backend, KVClassType.RECEIVER)
         kv_receiver = kv_receiver_class(
             mgr=self.kv_manager,
-            bootstrap_addr=f"{req.bootstrap_host}:{self.bootstrap_port}",
+            bootstrap_addr=f"{req.bootstrap_host}:{req.bootstrap_port}",
             bootstrap_room=req.bootstrap_room,
         )
         self.queue.append(DecodeRequest(req=req, kv_receiver=kv_receiver))
@@ -306,7 +307,7 @@ class DecodeTransferQueue:
     def extend(self, req_conns) -> None:
         self.queue.extend(req_conns)
 
-    def pop_transferred(self) -> List[Req]:
+    def pop_transferred(self) -> List[DecodeRequest]:
         if not self.queue:
             return []
 
@@ -329,7 +330,7 @@ class DecodeTransferQueue:
                 assert len(decode_req.req.output_ids) == 0
                 assert decode_req.req.transferred_output_id is None
                 decode_req.req.transferred_output_id = output_id
-                transferred_reqs.append(decode_req.req)
+                transferred_reqs.append(decode_req)
                 indices_to_remove.add(i)
             elif poll in [
                 KVPoll.Bootstrapping,
@@ -443,8 +444,17 @@ class ScheduleBatchDisaggregationDecodeMixin:
 
 class SchedulerDisaggregationDecodeMixin:
 
+    def _prepare_idle_batch_and_run(self, batch, delay_process=False):
+        batch, _ = self.prepare_dp_attn_batch(batch)
+        result = None
+        if batch:
+            result = self.run_batch(batch)
+            if not delay_process:
+                self.process_batch_result(batch, result)
+        return batch, result
+
     @torch.no_grad()
-    def event_loop_normal_disagg_decode(self):
+    def event_loop_normal_disagg_decode(self: Scheduler):
         """A normal scheduler loop for decode worker in disaggregation mode."""
 
         while True:
@@ -455,14 +465,25 @@ class SchedulerDisaggregationDecodeMixin:
             batch = self.get_next_disagg_decode_batch_to_run()
             self.cur_batch = batch
 
+            prepare_dp_attn_flag = (
+                self.server_args.enable_dp_attention
+                or self.server_args.enable_sp_layernorm
+            )
+
             if batch:
                 # Generate fake extend output.
                 if batch.forward_mode.is_extend():
                     # Note: Logprobs should be handled on the prefill engine.
                     self.stream_output(batch.reqs, False)
+                    if prepare_dp_attn_flag:
+                        self._prepare_idle_batch_and_run(None)
                 else:
+                    if prepare_dp_attn_flag:
+                        self.prepare_dp_attn_batch(batch)
                     result = self.run_batch(batch)
                     self.process_batch_result(batch, result)
+            elif prepare_dp_attn_flag:
+                batch, _ = self._prepare_idle_batch_and_run(None)
 
             if batch is None and (
                 len(self.disagg_decode_transfer_queue.queue)
@@ -475,6 +496,69 @@ class SchedulerDisaggregationDecodeMixin:
 
             self.last_batch = batch
 
+    @torch.no_grad()
+    def event_loop_overlap_disagg_decode(self: Scheduler):
+        result_queue = deque()
+        self.last_batch: Optional[ScheduleBatch] = None
+        self.last_batch_in_queue = False  # last batch is modifed in-place, so we need another variable to track if it's extend
+
+        while True:
+            recv_reqs = self.recv_requests()
+            self.process_input_requests(recv_reqs)
+            # polling and allocating kv cache
+            self.process_decode_queue()
+            batch = self.get_next_disagg_decode_batch_to_run()
+            self.cur_batch = batch
+            last_batch_in_queue = False
+
+            prepare_dp_attn_flag = (
+                self.server_args.enable_dp_attention
+                or self.server_args.enable_sp_layernorm
+            )
+
+            if batch:
+                # Generate fake extend output.
+                if batch.forward_mode.is_extend():
+                    # Note: Logprobs should be handled on the prefill engine.
+                    self.stream_output(batch.reqs, False)
+                    if prepare_dp_attn_flag:
+                        batch_, result = self._prepare_idle_batch_and_run(
+                            None, delay_process=True
+                        )
+                        if batch_:
+                            result_queue.append((batch_.copy(), result))
+                            last_batch_in_queue = True
+                else:
+                    if prepare_dp_attn_flag:
+                        self.prepare_dp_attn_batch(batch)
+                    result = self.run_batch(batch)
+                    result_queue.append((batch.copy(), result))
+                    last_batch_in_queue = True
+            elif prepare_dp_attn_flag:
+                batch, result = self._prepare_idle_batch_and_run(
+                    None, delay_process=True
+                )
+                if batch:
+                    result_queue.append((batch.copy(), result))
+                    last_batch_in_queue = True
+
+            # Process the results of the previous batch but skip if the last batch is extend
+            if self.last_batch and self.last_batch_in_queue:
+                tmp_batch, tmp_result = result_queue.popleft()
+                self.process_batch_result(tmp_batch, tmp_result)
+
+            if batch is None and (
+                len(self.disagg_decode_transfer_queue.queue)
+                + len(self.disagg_decode_prealloc_queue.queue)
+                == 0
+            ):
+                # When the server is idle, do self-check and re-init some states
+                self.check_memory()
+                self.new_token_ratio = self.init_new_token_ratio
+
+            self.last_batch = batch
+            self.last_batch_in_queue = last_batch_in_queue
+
     def get_next_disagg_decode_batch_to_run(
         self: Scheduler,
     ) -> Optional[Tuple[ScheduleBatch, bool]]:
@@ -557,8 +641,15 @@ class SchedulerDisaggregationDecodeMixin:
 
     def process_decode_queue(self: Scheduler):
         req_conns = self.disagg_decode_prealloc_queue.pop_preallocated()
+
+        def _num_pre_alloc(req):
+            return len(req.req.origin_input_ids) + max(len(req.req.output_ids) - 1, 0)
+
+        self.num_tokens_pre_allocated += sum(_num_pre_alloc(req) for req in req_conns)
         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)
+        self.num_tokens_pre_allocated -= sum(_num_pre_alloc(req) for req in alloc_reqs)
+
+        self.waiting_queue.extend([req.req for req in alloc_reqs])

sglang/srt/disaggregation/mini_lb.py

@@ -6,6 +6,7 @@ import asyncio
 import random
 import urllib
 from itertools import chain
+from typing import List
 
 import aiohttp
 import orjson
@@ -14,17 +15,27 @@ from fastapi import FastAPI, HTTPException
 from fastapi.responses import ORJSONResponse, Response, StreamingResponse
 
 
+class PrefillConfig:
+    def __init__(self, url: str, bootstrap_port: int):
+        self.url = url
+        self.bootstrap_port = bootstrap_port
+
+
 class MiniLoadBalancer:
-    def __init__(self, prefill_servers, decode_servers):
-        self.prefill_servers = prefill_servers
+    def __init__(self, prefill_configs: List[PrefillConfig], decode_servers: List[str]):
+        self.prefill_configs = prefill_configs
+        self.prefill_servers = [p.url for p in prefill_configs]
         self.decode_servers = decode_servers
 
     def select_pair(self):
-        return random.choice(self.prefill_servers), random.choice(self.decode_servers)
+        prefill_config = random.choice(self.prefill_configs)
+        decode_server = random.choice(self.decode_servers)
+        return prefill_config.url, prefill_config.bootstrap_port, decode_server
 
     async def generate(
-        self, modified_request, prefill_server, decode_server
+        self, modified_request, prefill_server, decode_server, endpoint
     ) -> ORJSONResponse:
+        assert endpoint[0] != "/", f"Endpoint should not start with '/': {endpoint}"
 
         async with aiohttp.ClientSession(
             timeout=aiohttp.ClientTimeout(
@@ -32,8 +43,8 @@ class MiniLoadBalancer:
             )  # Add timeout for request reliability
         ) as session:
             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. Prefill should end first.
             prefill_response, decode_response = await asyncio.gather(*tasks)
@@ -43,7 +54,11 @@ class MiniLoadBalancer:
                 status_code=decode_response.status,
             )
 
-    async def generate_stream(self, modified_request, prefill_server, decode_server):
+    async def generate_stream(
+        self, modified_request, prefill_server, decode_server, endpoint="generate"
+    ):
+        assert endpoint[0] != "/", f"Endpoint should not start with '/': {endpoint}"
+
         async def stream_results():
             async with aiohttp.ClientSession(
                 timeout=aiohttp.ClientTimeout(
@@ -54,10 +69,10 @@ class MiniLoadBalancer:
                     # Create the tasks for both prefill and decode requests
                     tasks = [
                         session.post(
-                            f"{prefill_server}/generate", json=modified_request
+                            f"{prefill_server}/{endpoint}", json=modified_request
                         ),
                         session.post(
-                            f"{decode_server}/generate", json=modified_request
+                            f"{decode_server}/{endpoint}", json=modified_request
                         ),
                     ]
                     # Wait for both responses to complete. Since this is streaming, they return immediately.
@@ -155,7 +170,46 @@ async def get_model_info():
 
 @app.post("/generate")
 async def handle_generate_request(request_data: dict):
-    prefill_server, decode_server = load_balancer.select_pair()
+    prefill_server, bootstrap_port, decode_server = load_balancer.select_pair()
+
+    # Parse and transform prefill_server for bootstrap data
+    parsed_url = urllib.parse.urlparse(prefill_server)
+    hostname = parsed_url.hostname
+    modified_request = request_data.copy()
+
+    batch_size = _get_request_batch_size(modified_request)
+    if batch_size is not None:
+        modified_request.update(
+            {
+                "bootstrap_host": [hostname] * batch_size,
+                "bootstrap_port": [bootstrap_port] * batch_size,
+                "bootstrap_room": [
+                    _generate_bootstrap_room() for _ in range(batch_size)
+                ],
+            }
+        )
+    else:
+        modified_request.update(
+            {
+                "bootstrap_host": hostname,
+                "bootstrap_port": bootstrap_port,
+                "bootstrap_room": _generate_bootstrap_room(),
+            }
+        )
+
+    if request_data.get("stream", False):
+        return await load_balancer.generate_stream(
+            modified_request, prefill_server, decode_server, "generate"
+        )
+    else:
+        return await load_balancer.generate(
+            modified_request, prefill_server, decode_server, "generate"
+        )
+
+
+@app.post("/v1/chat/completions")
+async def handle_completion_request(request_data: dict):
+    prefill_server, bootstrap_port, decode_server = load_balancer.select_pair()
 
     # Parse and transform prefill_server for bootstrap data
     parsed_url = urllib.parse.urlparse(prefill_server)
@@ -164,20 +218,40 @@ async def handle_generate_request(request_data: dict):
     modified_request.update(
         {
             "bootstrap_host": hostname,
+            "bootstrap_port": bootstrap_port,
             "bootstrap_room": random.randint(0, 2**63 - 1),
         }
     )
 
     if request_data.get("stream", False):
         return await load_balancer.generate_stream(
-            modified_request, prefill_server, decode_server
+            modified_request,
+            prefill_server,
+            decode_server,
+            endpoint="v1/chat/completions",
         )
     else:
         return await load_balancer.generate(
-            modified_request, prefill_server, decode_server
+            modified_request,
+            prefill_server,
+            decode_server,
+            endpoint="v1/chat/completions",
         )
 
 
+def _generate_bootstrap_room():
+    return random.randint(0, 2**63 - 1)
+
+
+# We may utilize `GenerateReqInput`'s logic later
+def _get_request_batch_size(request):
+    if (text := request.get("text")) is not None:
+        return None if isinstance(text, str) else len(text)
+    if (input_ids := request.get("input_ids")) is not None:
+        return None if isinstance(input_ids[0], int) else len(input_ids)
+    return None
+
+
 @app.get("/v1/models")
 async def get_models():
     prefill_server = load_balancer.prefill_servers[0]  # Get the first prefill server
@@ -194,9 +268,9 @@ async def get_models():
             raise HTTPException(status_code=500, detail=str(e))
 
 
-def run(prefill_addrs, decode_addrs, host, port):
+def run(prefill_configs, decode_addrs, host, port):
     global load_balancer
-    load_balancer = MiniLoadBalancer(prefill_addrs, decode_addrs)
+    load_balancer = MiniLoadBalancer(prefill_configs, decode_addrs)
     uvicorn.run(app, host=host, port=port)
 
 
@@ -207,6 +281,11 @@ if __name__ == "__main__":
     parser.add_argument(
         "--prefill", required=True, help="Comma-separated URLs for prefill servers"
     )
+    parser.add_argument(
+        "--prefill-bootstrap-ports",
+        help="Comma-separated bootstrap ports for prefill servers",
+        default="8998",
+    )
     parser.add_argument(
         "--decode", required=True, help="Comma-separated URLs for decode servers"
     )
@@ -217,4 +296,23 @@ if __name__ == "__main__":
         "--port", type=int, default=8000, help="Port to bind the server (default: 8000)"
     )
     args = parser.parse_args()
-    run(args.prefill.split(","), args.decode.split(","), args.host, args.port)
+
+    prefill_urls = args.prefill.split(",")
+    bootstrap_ports = [int(p) for p in args.prefill_bootstrap_ports.split(",")]
+
+    if len(bootstrap_ports) == 1:
+        bootstrap_ports = bootstrap_ports * len(prefill_urls)
+    else:
+        if len(bootstrap_ports) != len(prefill_urls):
+            raise ValueError(
+                "Number of prefill URLs must match number of bootstrap ports"
+            )
+            exit(1)
+
+    prefill_configs = []
+    for url, port in zip(prefill_urls, bootstrap_ports):
+        prefill_configs.append(PrefillConfig(url, port))
+
+    decode_addrs = args.decode.split(",")
+
+    run(prefill_configs, decode_addrs, args.host, args.port)