sglang 0.4.4__py3-none-any.whl → 0.4.4.post2__py3-none-any.whl

sglang/srt/disaggregation/mini_lb.py

@@ -0,0 +1,285 @@
+"""
+Minimal HTTP load balancer for prefill and decode servers for testing purpose.
+"""
+
+import asyncio
+import random
+import urllib
+from itertools import chain
+
+import aiohttp
+import orjson
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import ORJSONResponse, Response, StreamingResponse
+
+
+class MiniLoadBalancer:
+    def __init__(self, prefill_servers, decode_servers):
+        self.prefill_servers = prefill_servers
+        self.decode_servers = decode_servers
+
+    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 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),
+            ]
+
+            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()
+
+
+app = FastAPI()
+load_balancer = None
+
+
+@app.get("/health")
+async def health_check():
+    return Response(status_code=200)
+
+
+@app.get("/health_generate")
+async def health_check():
+    prefill_servers, decode_servers = (
+        load_balancer.prefill_servers,
+        load_balancer.decode_servers,
+    )
+    async with aiohttp.ClientSession() as session:
+        # Create the tasks
+        tasks = []
+        for server in chain(prefill_servers, decode_servers):
+            tasks.append(session.post(f"{server}/health_generate"))
+        for i, response in enumerate(asyncio.as_completed(tasks)):
+            await response
+    return Response(status_code=200)
+
+
+@app.post("/flush_cache")
+async def flush_cache():
+    prefill_servers, decode_servers = (
+        load_balancer.prefill_servers,
+        load_balancer.decode_servers,
+    )
+    async with aiohttp.ClientSession() as session:
+        # Create the tasks
+        tasks = []
+        for server in chain(prefill_servers, decode_servers):
+            tasks.append(session.post(f"{server}/flush_cache"))
+        for i, response in enumerate(asyncio.as_completed(tasks)):
+            await response
+    return Response(status_code=200)
+
+
+@app.get("/get_server_info")
+async def get_server_info():
+    prefill_servers, decode_servers = (
+        load_balancer.prefill_servers,
+        load_balancer.decode_servers,
+    )
+    prefill_infos = []
+    decode_infos = []
+    async with aiohttp.ClientSession() as session:
+        for server in chain(prefill_servers):
+            server_info = await session.get(f"{server}/get_server_info")
+            prefill_infos.append(await server_info.json())
+        for server in chain(decode_servers):
+            server_info = await session.get(f"{server}/get_server_info")
+            decode_infos.append(await server_info.json())
+
+    return {"prefill": prefill_infos, "decode": decode_infos}
+
+
+@app.get("/get_model_info")
+async def get_model_info():
+    # Dummy model information
+    model_info = {
+        "model_path": "/path/to/dummy/model",
+        "tokenizer_path": "/path/to/dummy/tokenizer",
+        "is_generation": True,
+        "preferred_sampling_params": {"temperature": 0.7, "max_new_tokens": 128},
+    }
+    return ORJSONResponse(content=model_info)
+
+
+@app.post("/generate")
+async def handle_generate_request(request_data: dict):
+    prefill_server, 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()
+    modified_request.update(
+        {
+            "bootstrap_host": hostname,
+            "bootstrap_room": random.randint(0, 2**63 - 1),
+        }
+    )
+
+    # 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",
+        )
+
+    # Non-streaming case
+    result = await load_balancer.generate_request(request_data)
+    return ORJSONResponse(content=result)
+
+
+@app.get("/v1/models")
+async def get_models():
+    prefill_server = load_balancer.prefill_servers[0]  # Get the first prefill server
+    async with aiohttp.ClientSession() as session:
+        try:
+            response = await session.get(f"{prefill_server}/v1/models")
+            if response.status != 200:
+                raise HTTPException(
+                    status_code=response.status,
+                    detail=f"Prefill server error: Status {response.status}",
+                )
+            return ORJSONResponse(content=await response.json())
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
+
+
+def run(prefill_addrs, decode_addrs, host, port):
+    global load_balancer
+    load_balancer = MiniLoadBalancer(prefill_addrs, decode_addrs)
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Mini Load Balancer Server")
+    parser.add_argument(
+        "--prefill", required=True, help="Comma-separated URLs for prefill servers"
+    )
+    parser.add_argument(
+        "--decode", required=True, help="Comma-separated URLs for decode 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()
+    run(args.prefill.split(","), args.decode.split(","), args.host, args.port)

sglang/srt/disaggregation/prefill.py

@@ -0,0 +1,249 @@
+"""
+Life cycle of a request in the prefill server
+
+1. Bootstrap Queue
+    a. Initialize a sender for each request
+    b. Use the queue to store requests whose bootstrap (handshake and preallocation) has not finished
+    c. Poll senders to check bootstrap state
+    d. Once bootstrap is complete, move request to Waiting Queue
+
+2. Waiting Queue
+    a. Use PrefillAdder to pop requests
+    b. Run forward
+    c. Add the request to Infight Queue
+
+3. Infight Queue
+    a. Poll (non-blocking) the sender of the request
+    b. Once the transfer has finished, return the request
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+
+from sglang.srt.disaggregation.conn import KVArgs, KVManager, KVPoll, KVSender
+from sglang.srt.disaggregation.utils import (
+    ReqToMetadataIdxAllocator,
+    poll_and_all_reduce,
+)
+from sglang.srt.managers.schedule_batch import FINISH_LENGTH, Req, ScheduleBatch
+
+if TYPE_CHECKING:
+    from torch.distributed import ProcessGroup
+
+    from sglang.srt.managers.scheduler import GenerationBatchResult, Scheduler
+    from sglang.srt.mem_cache.memory_pool import KVCache
+
+logger = logging.getLogger(__name__)
+
+
+class PrefillBootstrapQueue:
+    """
+    Store the requests in bootstrapping
+    """
+
+    def __init__(
+        self,
+        token_to_kv_pool: KVCache,
+        req_to_metadata_buffer_idx_allocator: ReqToMetadataIdxAllocator,
+        metadata_buffers: List[torch.Tensor],
+        aux_dtype: torch.dtype,
+        tp_rank: int,
+        tp_size: int,
+        bootstrap_port: int,
+        gloo_group: ProcessGroup,
+    ):
+        self.token_to_kv_pool = token_to_kv_pool
+        self.aux_dtype = aux_dtype
+
+        self.metadata_buffers = metadata_buffers
+        self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator
+        self.tp_rank = tp_rank
+        self.tp_size = tp_size
+        self.kv_manager = self._init_kv_manager()
+        self.queue: List[Req] = []
+        self.gloo_group = gloo_group
+        self.bootstrap_port = bootstrap_port
+
+    def allocate_token_id(self, idx: int, token_id: int):
+        assert token_id >= 0, f"token_id: {token_id} is negative"
+        output_id_buffer = self.metadata_buffers[0]
+        output_id_buffer[idx] = token_id
+
+    def _init_kv_manager(self) -> KVManager:
+        kv_args = KVArgs()
+        kv_args.engine_rank = self.tp_rank
+        kv_data_ptrs, kv_data_lens, kv_item_lens = (
+            self.token_to_kv_pool.get_contiguous_buf_infos()
+        )
+
+        kv_args.kv_data_ptrs = kv_data_ptrs
+        kv_args.kv_data_lens = kv_data_lens
+        kv_args.kv_item_lens = kv_item_lens
+
+        # Define req -> input ids buffer
+        kv_args.aux_data_ptrs = [
+            metadata_buffer.data_ptr() for metadata_buffer in self.metadata_buffers
+        ]
+        kv_args.aux_data_lens = [
+            metadata_buffer.nbytes for metadata_buffer in self.metadata_buffers
+        ]
+        kv_args.aux_item_lens = [
+            metadata_buffer[0].nbytes for metadata_buffer in self.metadata_buffers
+        ]
+        kv_args.ib_device = "mock-ib-device"
+        kv_manager = KVManager(kv_args)
+        return kv_manager
+
+    def add(self, req: Req) -> None:
+        req.disagg_kv_sender = KVSender(
+            mgr=self.kv_manager,
+            bootstrap_addr=f"{req.bootstrap_host}:{self.bootstrap_port}",
+            bootstrap_room=req.bootstrap_room,
+        )
+        self._process_req(req)
+        self.queue.append(req)
+
+    def _process_req(self, req: Req) -> None:
+        """
+        Set max_new_tokens = 1, so PrefillAdder memory estimation is accurate
+        """
+        req.sampling_params.max_new_tokens = 1
+
+    def pop_bootstrapped(self) -> List[Req]:
+        """pop the reqs which has finished bootstrapping"""
+        bootstrapped_reqs = []
+        indices_to_remove = set()
+
+        if len(self.queue) == 0:
+            return []
+
+        polls = poll_and_all_reduce(
+            [req.disagg_kv_sender for req in self.queue], self.gloo_group
+        )
+
+        for i, (req, poll) in enumerate(zip(self.queue, polls)):
+            if poll == KVPoll.Bootstrapping:
+                continue
+            elif poll == KVPoll.Failed:
+                raise Exception("Bootstrap failed")
+
+            # KV.WaitingForInput - init here
+            num_kv_indices = len(req.origin_input_ids)
+            if self.req_to_metadata_buffer_idx_allocator.available_size() == 0:
+                break
+
+            req.metadata_buffer_index = (
+                self.req_to_metadata_buffer_idx_allocator.alloc()
+            )
+            assert req.metadata_buffer_index is not None
+            req.disagg_kv_sender.init(num_kv_indices, req.metadata_buffer_index)
+
+            bootstrapped_reqs.append(req)
+            indices_to_remove.add(i)
+
+        self.queue = [
+            entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
+        ]
+
+        return bootstrapped_reqs
+
+
+class SchedulerDisaggregationPrefillMixin:
+    """
+    Mixin for Scheduler to handle disaggregation prefill
+    """
+
+    def process_batch_result_disagg_prefill(
+        self: Scheduler, batch: ScheduleBatch, result: GenerationBatchResult
+    ) -> None:
+        """
+        Transfer kv for prefill completed requests and add it into disagg_prefill_infight_queue
+        Adapted from process_batch_result_prefill
+        """
+
+        next_token_ids = result.next_token_ids.tolist()
+
+        for req, next_token_id in zip(batch.reqs, next_token_ids, strict=True):
+            req: Req
+            if req.is_chunked <= 0:
+                # There is no output_ids for prefill
+                req.output_ids.append(next_token_id)
+                self.tree_cache.cache_unfinished_req(req)  # update the tree and lock
+                self.send_kv_chunk(req, token_id=next_token_id)
+                self.disagg_prefill_infight_queue.append(req)
+            else:
+                # being chunked reqs' prefill is not finished
+                req.is_chunked -= 1
+
+        # TODO: Not sure if this is necessary
+        if batch.next_batch_sampling_info:
+            batch.next_batch_sampling_info.update_regex_vocab_mask()
+            # We need to remove this for overlap schedule.
+            self.current_stream.synchronize()
+            batch.next_batch_sampling_info.sampling_info_done.set()
+
+    def process_disagg_prefill_infight_queue(self: Scheduler) -> None:
+        """
+        Poll the requests in the middle of transfer. If done, return the request.
+        """
+        assert len(self.disagg_prefill_infight_queue) > 0
+
+        done_reqs = []
+
+        polls = poll_and_all_reduce(
+            [req.disagg_kv_sender for req in self.disagg_prefill_infight_queue],
+            self.tp_worker.get_tp_cpu_group(),
+        )
+
+        undone_reqs: List[Req] = []
+        # Check .poll() for the reqs in disagg_prefill_infight_queue. If Success, respond to the client and remove it from the queue
+        for req, poll in zip(self.disagg_prefill_infight_queue, polls):
+            if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]:
+                undone_reqs.append(req)
+            elif poll == KVPoll.Success:  # transfer done
+                self.tree_cache.cache_finished_req(req)  # unlock the tree
+                req.finished_reason = FINISH_LENGTH(length=0)
+                done_reqs.append(req)
+            elif poll == KVPoll.Failed:
+                raise Exception("Transferring failed")
+
+        # Stream requests which have finished transfer
+        self.stream_output(done_reqs, False, None)
+
+        self.disagg_prefill_infight_queue = undone_reqs
+
+    def process_prefill_chunk(self: Scheduler) -> None:
+        if self.last_batch and self.last_batch.forward_mode.is_extend():
+            if self.chunked_req:
+                # Move the chunked request out of the batch so that we can merge
+                # only finished requests to running_batch.
+                self.last_batch.filter_batch(chunked_req_to_exclude=self.chunked_req)
+                self.tree_cache.cache_unfinished_req(self.chunked_req)
+                self.send_kv_chunk(self.chunked_req)
+                # chunked request keeps its rid but will get a new req_pool_idx
+                self.req_to_token_pool.free(self.chunked_req.req_pool_idx)
+                self.running_batch.batch_is_full = False
+
+    def send_kv_chunk(
+        self: Scheduler, req: Req, token_id: Optional[int] = None
+    ) -> None:
+        """
+        Send a prefilled chunk to the decode server
+        """
+        start_idx = req.start_send_idx
+        end_idx = min(len(req.fill_ids), len(req.origin_input_ids))
+        kv_indices = (
+            self.req_to_token_pool.req_to_token[req.req_pool_idx][start_idx:end_idx]
+            .cpu()
+            .numpy()
+        )
+        req.start_send_idx = end_idx
+        if token_id is not None:
+            self.disagg_prefill_pending_queue.allocate_token_id(
+                req.metadata_buffer_index, token_id
+            )
+        req.disagg_kv_sender.send(kv_indices)

sglang/srt/disaggregation/utils.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from collections import deque
+from enum import Enum
+from typing import List
+
+import torch
+import torch.distributed as dist
+
+
+class DisaggregationMode(Enum):
+    NULL = "null"
+    PREFILL = "prefill"
+    DECODE = "decode"
+
+
+def poll_and_all_reduce(pollers, gloo_group):
+    polls = [int(poller.poll()) for poller in pollers]
+    tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu")
+    dist.all_reduce(tensor_to_reduce, op=dist.ReduceOp.MIN, group=gloo_group)
+    return tensor_to_reduce.tolist()
+
+
+class ReqToMetadataIdxAllocator:
+    """A memory pool that maps a request to its first output token location."""
+
+    def __init__(
+        self,
+        size: int,
+    ):
+        self.size = size
+        self.free_slots = deque(list(range(size)))
+
+    def available_size(self):
+        return len(self.free_slots)
+
+    def alloc(self) -> List[int]:
+        if len(self.free_slots) == 0:
+            return None
+
+        return self.free_slots.popleft()
+
+    def free(self, free_index: int):
+        self.free_slots.append(free_index)

sglang/srt/distributed/parallel_state.py

@@ -189,6 +189,9 @@ class GroupCoordinator:
     device_group: ProcessGroup  # group for device communication
     use_pynccl: bool  # a hint of whether to use PyNccl
     use_custom_allreduce: bool  # a hint of whether to use CustomAllreduce
+    use_message_queue_broadcaster: (
+        bool  # a hint of whether to use message queue broadcaster
+    )
     # communicators are only created for world size > 1
     pynccl_comm: Optional[Any]  # PyNccl communicator
     ca_comm: Optional[Any]  # Custom allreduce communicator
@@ -241,6 +244,7 @@ class GroupCoordinator:
         self.use_custom_allreduce = use_custom_allreduce
         self.use_hpu_communicator = use_hpu_communicator
         self.use_xpu_communicator = use_xpu_communicator
+        self.use_message_queue_broadcaster = use_message_queue_broadcaster
 
         # lazy import to avoid documentation build error
         from sglang.srt.distributed.device_communicators.custom_all_reduce import (
@@ -269,7 +273,7 @@ class GroupCoordinator:
             HpuCommunicator,
         )
 
-        self.hpu_communicator: Optional[HpuCommunicator]
+        self.hpu_communicator: Optional[HpuCommunicator] = None
         if use_hpu_communicator and self.world_size > 1:
             self.hpu_communicator = HpuCommunicator(group=self.device_group)
 
@@ -277,7 +281,7 @@ class GroupCoordinator:
             XpuCommunicator,
         )
 
-        self.xpu_communicator: Optional[XpuCommunicator]
+        self.xpu_communicator: Optional[XpuCommunicator] = None
         if use_xpu_communicator and self.world_size > 1:
             self.xpu_communicator = XpuCommunicator(group=self.device_group)
 
@@ -1312,7 +1316,10 @@ vllm_get_world_group = None
 
 
 def monkey_patch_vllm_parallel_state(reverse: bool = False):
-    import vllm.distributed.parallel_state as vllm_parrlel_state
+    try:
+        import vllm.distributed.parallel_state as vllm_parrlel_state
+    except ImportError:
+        return
 
     global vllm_get_pp_group, vllm_get_tp_group, vllm_get_world_group
     if vllm_get_pp_group is None: