sglang 0.4.4.post4__py3-none-any.whl → 0.4.5.post1__py3-none-any.whl

sglang/srt/disaggregation/decode.py

@@ -24,12 +24,17 @@ import logging
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
+import numpy as np
 import torch
 from torch.distributed import ProcessGroup
 
-from sglang.srt.disaggregation.conn import KVArgs, KVManager, KVPoll, KVReceiver
+from sglang.srt.disaggregation.base import BaseKVManager, BaseKVReceiver, KVArgs, KVPoll
 from sglang.srt.disaggregation.utils import (
+    DisaggregationMode,
+    KVClassType,
     ReqToMetadataIdxAllocator,
+    TransferBackend,
+    get_kv_class,
     poll_and_all_reduce,
 )
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
@@ -49,7 +54,7 @@ if TYPE_CHECKING:
 @dataclass
 class DecodeRequest:
     req: Req
-    kv_receiver: KVReceiver
+    kv_receiver: BaseKVReceiver
     waiting_for_input: bool = False
     metadata_buffer_index: int = -1
 
@@ -73,6 +78,7 @@ class DecodePreallocQueue:
         tp_rank: int,
         tp_size: int,
         bootstrap_port: int,
+        transfer_backend: TransferBackend,
     ):
         self.req_to_token_pool = req_to_token_pool
         self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
@@ -92,9 +98,10 @@ class DecodePreallocQueue:
 
         # Queue for requests pending pre-allocation
         self.queue: List[DecodeRequest] = []
+        self.transfer_backend = transfer_backend
         self.kv_manager = self._init_kv_manager()
 
-    def _init_kv_manager(self) -> KVManager:
+    def _init_kv_manager(self) -> BaseKVManager:
         kv_args = KVArgs()
         kv_args.engine_rank = self.tp_rank
         kv_data_ptrs, kv_data_lens, kv_item_lens = (
@@ -115,13 +122,18 @@ class DecodePreallocQueue:
             metadata_buffer[0].nbytes for metadata_buffer in self.metadata_buffers
         ]
         kv_args.ib_device = "mock-ib-device"
-        kv_manager = KVManager(kv_args)
+        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_args, DisaggregationMode.DECODE, self.scheduler.server_args
+        )
         return kv_manager
 
     def add(self, req: Req) -> None:
         """Add a request to the pending queue."""
 
-        kv_receiver = KVReceiver(
+        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_room=req.bootstrap_room,
@@ -186,6 +198,7 @@ class DecodePreallocQueue:
                 ]
                 .cpu()
                 .numpy()
+                .astype(np.int64)
             )
 
             decode_req.metadata_buffer_index = (

sglang/srt/disaggregation/mini_lb.py

@@ -1,5 +1,5 @@
 """
-Minimal HTTP load balancer for prefill and decode servers for testing purpose.
+Minimal HTTP load balancer for prefill and decode servers for testing.
 """
 
 import asyncio
@@ -22,64 +22,59 @@ class MiniLoadBalancer:
     def select_pair(self):
         return random.choice(self.prefill_servers), random.choice(self.decode_servers)
 
-    async def generate_request(self, request_data):
-        prefill_server, decode_server = self.select_pair()
-
-        # Parse and transform prefill_server
-        parsed_url = urllib.parse.urlparse(prefill_server)
-        hostname = parsed_url.hostname
-        bootstrap_host = f"{hostname}"
-
-        modified_request = request_data.copy()
-        modified_request.update(
-            {
-                "bootstrap_host": bootstrap_host,
-                "bootstrap_room": random.randint(0, 2**63 - 1),
-            }
-        )
+    async def generate(
+        self, modified_request, prefill_server, decode_server
+    ) -> ORJSONResponse:
 
         async with aiohttp.ClientSession() as session:
-            # Create the tasks
             tasks = [
                 session.post(f"{prefill_server}/generate", json=modified_request),
                 session.post(f"{decode_server}/generate", json=modified_request),
             ]
+            # Wait for both responses to complete. Prefill should end first.
+            prefill_response, decode_response = await asyncio.gather(*tasks)
+
+            return ORJSONResponse(
+                content=await decode_response.json(),
+                status_code=decode_response.status,
+            )
+
+    async def generate_stream(self, modified_request, prefill_server, decode_server):
+        async def stream_results():
+            async with aiohttp.ClientSession(
+                timeout=aiohttp.ClientTimeout(
+                    total=3600
+                )  # Add timeout for request reliability
+            ) as session:
+                try:
+                    # 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
+                        ),
+                    ]
+                    # Wait for both responses to complete. Since this is streaming, they return immediately.
+                    prefill_response, decode_response = await asyncio.gather(*tasks)
+                    async for chunk in decode_response.content:
+                        yield chunk
+                except Exception as e:
+                    error_msg = {
+                        "error": {"message": f"Stream processing error: {str(e)}"}
+                    }
+                    yield b"data: " + orjson.dumps(
+                        error_msg, option=orjson.OPT_NON_STR_KEYS
+                    ) + b"\n\n"
+                finally:
+                    if prefill_response is not None:
+                        await prefill_response.release()
 
-            prefill_response = None
-            decode_response = None
-
-            # Process responses as they arrive
-            for i, response in enumerate(asyncio.as_completed(tasks)):
-                response = await response
-                # Check if this is the prefill or decode response based on order created
-                if i == 0:  # First completed task
-                    if str(response.url).startswith(prefill_server):
-                        prefill_response = response
-                        if response.status != 200:
-                            raise HTTPException(
-                                status_code=response.status,
-                                detail=f"Prefill server error: Status {response.status} Details: {await response.text()}",
-                            )
-                    else:
-                        decode_response = response
-                        if response.status != 200:
-                            raise HTTPException(
-                                status_code=response.status,
-                                detail=f"Decode server error: Status {response.status} Details: {await response.text()}",
-                            )
-                else:  # Second completed task
-                    if str(response.url).startswith(prefill_server):
-                        prefill_response = response
-                    else:
-                        decode_response = response
-
-                    if response.status != 200:
-                        raise HTTPException(
-                            status_code=response.status,
-                            detail=f"{'Prefill' if str(response.url).startswith(prefill_server) else 'Decode'} server error: Status {response.status} Details: {await response.text()}",
-                        )
-
-            return await decode_response.json()
+        return StreamingResponse(
+            stream_results(),
+            media_type="text/event-stream",
+        )
 
 
 app = FastAPI()
@@ -169,78 +164,14 @@ async def handle_generate_request(request_data: dict):
         }
     )
 
-    # Check if streaming is requested
     if request_data.get("stream", False):
-
-        async def stream_results():
-            async with aiohttp.ClientSession(
-                timeout=aiohttp.ClientTimeout(total=3600)
-            ) as session:
-                try:
-                    # Create the tasks
-                    tasks = [
-                        session.post(
-                            f"{prefill_server}/generate", json=modified_request
-                        ),
-                        session.post(
-                            f"{decode_server}/generate", json=modified_request
-                        ),
-                    ]
-
-                    prefill_response = None
-                    decode_response = None
-
-                    # Process responses as they arrive
-                    for i, response_task in enumerate(asyncio.as_completed(tasks)):
-                        response = await response_task
-
-                        # Check the response immediately
-                        if str(response.url).startswith(prefill_server):
-                            prefill_response = response
-                            if response.status != 200:
-                                error_msg = {
-                                    "error": {
-                                        "message": f"Prefill server error: Status {response.status}, Details: {await response.text()}"
-                                    }
-                                }
-                                yield b"data: " + orjson.dumps(
-                                    error_msg, option=orjson.OPT_NON_STR_KEYS
-                                ) + b"\n\n"
-                                return
-                        else:
-                            decode_response = response
-                            if response.status != 200:
-                                error_msg = {
-                                    "error": {
-                                        "message": f"Decode server error: Status {response.status}"
-                                    }
-                                }
-                                yield b"data: " + orjson.dumps(
-                                    error_msg, option=orjson.OPT_NON_STR_KEYS
-                                ) + b"\n\n"
-                                return
-
-                    # Stream successful decode server response
-                    async for line in decode_response.content:
-                        yield line
-                    yield b"data: [DONE]\n\n"
-
-                except Exception as e:
-                    error_msg = {
-                        "error": {"message": f"Stream processing error: {str(e)}"}
-                    }
-                    yield b"data: " + orjson.dumps(
-                        error_msg, option=orjson.OPT_NON_STR_KEYS
-                    ) + b"\n\n"
-
-        return StreamingResponse(
-            stream_results(),
-            media_type="text/event-stream",
+        return await load_balancer.generate_stream(
+            modified_request, prefill_server, decode_server
+        )
+    else:
+        return await load_balancer.generate(
+            modified_request, prefill_server, decode_server
         )
-
-    # Non-streaming case
-    result = await load_balancer.generate_request(request_data)
-    return ORJSONResponse(content=result)
 
 
 @app.get("/v1/models")

sglang/srt/disaggregation/mooncake/__init__.py

@@ -0,0 +1,6 @@
+from .conn import (
+    MooncakeKVBootstrapServer,
+    MooncakeKVManager,
+    MooncakeKVReceiver,
+    MooncakeKVSender,
+)