sglang 0.4.5.post1__py3-none-any.whl → 0.4.5.post2__py3-none-any.whl

sglang/srt/disaggregation/mooncake/conn.py

@@ -99,8 +99,12 @@ class MooncakeKVManager(BaseKVManager):
         disaggregation_mode: DisaggregationMode,
         server_args: ServerArgs,
     ):
-        self.engine = MooncakeTransferEngine()
         self.kv_args = args
+        self.engine = MooncakeTransferEngine(
+            hostname=get_local_ip_by_remote(),
+            gpu_id=self.kv_args.gpu_id,
+            ib_device=self.kv_args.ib_device,
+        )
         self.disaggregation_mode = disaggregation_mode
         # for p/d multi node infer
         self.bootstrap_port = server_args.disaggregation_bootstrap_port
@@ -387,6 +391,10 @@ class MooncakeKVSender(BaseKVSender):
 
 
 class MooncakeKVReceiver(BaseKVReceiver):
+    _ctx = zmq.Context()
+    _socket_cache = {}
+    _socket_locks = {}
+    _global_lock = threading.Lock()
 
     def __init__(
         self,
@@ -436,11 +444,15 @@ class MooncakeKVReceiver(BaseKVReceiver):
             logger.error(f"Error fetching prefill info from bootstrap: {e}")
             return None
 
-    @cache
-    def _connect(self, endpoint: str):
-        socket = zmq.Context().socket(zmq.PUSH)
-        socket.connect(endpoint)
-        return socket
+    @classmethod
+    def _connect(cls, endpoint: str):
+        with cls._global_lock:
+            if endpoint not in cls._socket_cache:
+                sock = cls._ctx.socket(zmq.PUSH)
+                sock.connect(endpoint)
+                cls._socket_cache[endpoint] = sock
+                cls._socket_locks[endpoint] = threading.Lock()
+            return cls._socket_cache[endpoint], cls._socket_locks[endpoint]
 
     def init(self, kv_indices: npt.NDArray[np.int64], aux_index: Optional[int] = None):
         self.prefill_server_url = (
@@ -456,18 +468,20 @@ class MooncakeKVReceiver(BaseKVReceiver):
         packed_aux_data_ptrs = b"".join(
             struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.aux_data_ptrs
         )
-        self._connect("tcp://" + self.prefill_server_url).send_multipart(
-            [
-                str(self.bootstrap_room).encode("ascii"),
-                get_local_ip_by_remote().encode("ascii"),
-                str(self.kv_mgr.rank_port).encode("ascii"),
-                self.session_id.encode("ascii"),
-                packed_kv_data_ptrs,
-                kv_indices.tobytes(),
-                packed_aux_data_ptrs,
-                str(aux_index).encode("ascii"),
-            ]
-        )
+        sock, lock = self._connect("tcp://" + self.prefill_server_url)
+        with lock:
+            sock.send_multipart(
+                [
+                    str(self.bootstrap_room).encode("ascii"),
+                    get_local_ip_by_remote().encode("ascii"),
+                    str(self.kv_mgr.rank_port).encode("ascii"),
+                    self.session_id.encode("ascii"),
+                    packed_kv_data_ptrs,
+                    kv_indices.tobytes(),
+                    packed_aux_data_ptrs,
+                    str(aux_index).encode("ascii"),
+                ]
+            )
 
     def poll(self) -> KVPoll:
         return self.kv_mgr.check_status(self.bootstrap_room)
@@ -493,52 +507,8 @@ class MooncakeKVBootstrapServer(BaseKVBootstrapServer):
         self.thread.start()
 
     def _setup_routes(self):
-        self.app.router.add_route("*", "/metadata", self._handle_metadata)
         self.app.router.add_route("*", "/route", self._handle_route)
 
-    async def _handle_metadata(self, request: web.Request):
-        key = request.query.get("key", "")
-
-        if request.method == "GET":
-            return await self._handle_metadata_get(key)
-        elif request.method == "PUT":
-            return await self._handle_metadata_put(key, request)
-        elif request.method == "DELETE":
-            return await self._handle_metadata_delete(key)
-        return web.Response(
-            text="Method not allowed", status=405, content_type="application/json"
-        )
-
-    async def _handle_metadata_get(self, key):
-        async with self.lock:
-            value = self.store.get(key)
-        if value is None:
-            return web.Response(
-                text="metadata not found", status=404, content_type="application/json"
-            )
-        return web.Response(body=value, status=200, content_type="application/json")
-
-    async def _handle_metadata_put(self, key, request):
-        data = await request.read()
-        async with self.lock:
-            self.store[key] = data
-        return web.Response(
-            text="metadata updated", status=200, content_type="application/json"
-        )
-
-    async def _handle_metadata_delete(self, key):
-        async with self.lock:
-            if key not in self.store:
-                return web.Response(
-                    text="metadata not found",
-                    status=404,
-                    content_type="application/json",
-                )
-            del self.store[key]
-        return web.Response(
-            text="metadata deleted", status=200, content_type="application/json"
-        )
-
     async def _handle_route(self, request: web.Request):
         method = request.method
         if method == "PUT":

sglang/srt/disaggregation/mooncake/transfer_engine.py

@@ -1,45 +1,14 @@
 import json
 import logging
-import os
-import uuid
 from dataclasses import dataclass
+from typing import Optional
 
 logger = logging.getLogger(__name__)
 
 
-@dataclass
-class MooncakeTransferEngineConfig:
-    local_hostname: str
-    metadata_server: str
-    protocol: str
-    device_name: str
-
-    @staticmethod
-    def from_file(file_path: str) -> "MooncakeTransferEngineConfig":
-        """Load the config from a JSON file."""
-        with open(file_path) as fin:
-            config = json.load(fin)
-        return MooncakeTransferEngineConfig(
-            local_hostname=config.get("local_hostname", None),
-            metadata_server=config.get("metadata_server"),
-            protocol=config.get("protocol", "rdma"),
-            device_name=config.get("device_name", ""),
-        )
-
-    @staticmethod
-    def load_from_env() -> "MooncakeTransferEngineConfig":
-        """Load config from a file specified in the environment variable."""
-        config_file_path = os.getenv("MOONCAKE_CONFIG_PATH")
-        if config_file_path is None:
-            raise ValueError(
-                "The environment variable 'MOONCAKE_CONFIG_PATH' is not set."
-            )
-        return MooncakeTransferEngineConfig.from_file(config_file_path)
-
-
 class MooncakeTransferEngine:
 
-    def __init__(self):
+    def __init__(self, hostname: str, gpu_id: int, ib_device: Optional[str] = None):
         try:
             from mooncake.engine import TransferEngine
         except ImportError as e:
@@ -50,43 +19,43 @@ class MooncakeTransferEngine:
             ) from e
 
         self.engine = TransferEngine()
+        self.hostname = hostname
+        self.gpu_id = gpu_id
+        self.ib_device = ib_device
 
-        try:
-            self.config = MooncakeTransferEngineConfig.load_from_env()
-            logger.info("Mooncake Configuration loaded successfully.")
-        except ValueError as e:
-            logger.error(e)
-            raise
-        except Exception as exc:
-            logger.error("An error occurred while loading the configuration: %s", exc)
-            raise
-
-        self.config = MooncakeTransferEngineConfig.load_from_env()
-
-        session_suffix = "_" + str(uuid.uuid4())
-        self.session_id = self.config.local_hostname + session_suffix
         self.initialize(
-            self.session_id,
-            self.config.metadata_server,
-            self.config.protocol,
-            self.config.device_name,
+            hostname=self.hostname,
+            device_name=self.ib_device,
         )
+        self.session_id = f"{self.hostname}:{self.engine.get_rpc_port()}"
 
     def register(self, ptr, length):
-        self.engine.register_memory(ptr, length)
+        ret_value = self.engine.register_memory(ptr, length)
+        if ret_value != 0:
+            logger.error("Mooncake memory registration failed.")
+            raise RuntimeError("Mooncake memory registration failed.")
 
     def deregister(self, ptr):
-        self.engine.unregister_memory(ptr)
+        ret_value = self.engine.unregister_memory(ptr)
+        if ret_value != 0:
+            logger.error("Mooncake memory deregistration failed.")
+            raise RuntimeError("Mooncake memory deregistration failed.")
 
     def initialize(
         self,
-        local_hostname: str,
-        metadata_server: str,
-        protocol: str,
-        device_name: str,
+        hostname: str,
+        device_name: Optional[str],
     ) -> None:
         """Initialize the mooncake instance."""
-        self.engine.initialize(local_hostname, metadata_server, protocol, device_name)
+        ret_value = self.engine.initialize(
+            hostname,
+            "P2PHANDSHAKE",
+            "rdma",
+            device_name if device_name is not None else "",
+        )
+        if ret_value != 0:
+            logger.error("Mooncake Transfer Engine initialization failed.")
+            raise RuntimeError("Mooncake Transfer Engine initialization failed.")
 
     def transfer_sync(
         self, session_id: str, buffer: int, peer_buffer_address: int, length: int
@@ -97,12 +66,12 @@ class MooncakeTransferEngine:
             session_id, buffer, peer_buffer_address, length
         )
         if ret < 0:
-            logger.error("Transfer Return Error")
-            raise Exception("Transfer Return Error")
+            logger.error("Mooncake Transfer Engine Return Error.")
+            raise RuntimeError("Mooncake Transfer Engine Return Error.")
         return ret
 
     def get_localhost(self):
-        return self.config.local_hostname
+        return self.hostname
 
     def get_session_id(self):
         return self.session_id

sglang/srt/disaggregation/prefill.py

@@ -31,6 +31,8 @@ from sglang.srt.disaggregation.utils import (
     ReqToMetadataIdxAllocator,
     TransferBackend,
     get_kv_class,
+    kv_to_page_indices,
+    kv_to_page_num,
     poll_and_all_reduce,
 )
 from sglang.srt.managers.schedule_batch import FINISH_LENGTH, Req, ScheduleBatch
@@ -103,7 +105,7 @@ class PrefillBootstrapQueue:
         kv_args.aux_item_lens = [
             metadata_buffer[0].nbytes for metadata_buffer in self.metadata_buffers
         ]
-        kv_args.ib_device = "mock-ib-device"
+        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(
@@ -154,7 +156,8 @@ class PrefillBootstrapQueue:
                 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)
+            num_pages = kv_to_page_num(num_kv_indices, self.token_to_kv_pool.page_size)
+            req.disagg_kv_sender.init(num_pages, req.metadata_buffer_index)
 
             bootstrapped_reqs.append(req)
             indices_to_remove.add(i)
@@ -171,6 +174,36 @@ class SchedulerDisaggregationPrefillMixin:
     Mixin for Scheduler to handle disaggregation prefill
     """
 
+    @torch.no_grad()
+    def event_loop_normal_disagg_prefill(self):
+        """A normal scheduler loop for prefill worker in disaggregation mode."""
+
+        while True:
+            recv_reqs = self.recv_requests()
+            self.process_input_requests(recv_reqs)
+            self.waiting_queue.extend(
+                self.disagg_prefill_pending_queue.pop_bootstrapped()
+            )
+            self.process_prefill_chunk()
+            batch = self.get_new_batch_prefill()
+            self.cur_batch = batch
+
+            if batch:
+                result = self.run_batch(batch)
+                self.process_batch_result_disagg_prefill(batch, result)
+
+            if len(self.disagg_prefill_inflight_queue) > 0:
+                self.process_disagg_prefill_inflight_queue()
+
+            if batch is None and len(self.disagg_prefill_inflight_queue) == 0:
+                self.check_memory()
+                self.new_token_ratio = self.init_new_token_ratio
+
+            self.last_batch = batch
+            # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it
+            # Otherwise, it hangs under high concurrency
+            self.running_batch.batch_is_full = False
+
     def process_batch_result_disagg_prefill(
         self: Scheduler, batch: ScheduleBatch, result: GenerationBatchResult
     ) -> None:
@@ -210,7 +243,7 @@ class SchedulerDisaggregationPrefillMixin:
 
         polls = poll_and_all_reduce(
             [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue],
-            self.tp_worker.get_tp_cpu_group(),
+            self.attn_tp_cpu_group,
         )
 
         undone_reqs: List[Req] = []
@@ -270,4 +303,7 @@ class SchedulerDisaggregationPrefillMixin:
                 req.metadata_buffer_index, token_id
             )
         is_last = token_id is not None
-        req.disagg_kv_sender.send(kv_indices, slice(start_idx, end_idx), is_last)
+        page_indices = kv_to_page_indices(
+            kv_indices, self.token_to_kv_pool_allocator.page_size
+        )
+        req.disagg_kv_sender.send(page_indices, slice(start_idx, end_idx), is_last)

sglang/srt/disaggregation/utils.py

@@ -4,6 +4,7 @@ from collections import deque
 from enum import Enum
 from typing import List
 
+import numpy as np
 import torch
 import torch.distributed as dist
 
@@ -73,3 +74,17 @@ def get_kv_class(transfer_backend: TransferBackend, class_type: KVClassType):
         }
         return class_mapping.get(class_type)
     raise ValueError(f"Unsupported transfer backend: {transfer_backend}")
+
+
+def kv_to_page_indices(kv_indices: np.ndarray, page_size: int):
+    # 1. The page is guaruanteed to be full except the last page.
+    # 2. page index = kv_index // page_size
+    # The return vector is kv_indices[::page_size] // page_size
+    if page_size == 1:  # shortcut
+        return kv_indices
+    return kv_indices[::page_size] // page_size
+
+
+def kv_to_page_num(num_kv_indices: int, page_size: int):
+    # ceil(num_kv_indices / page_size)
+    return (num_kv_indices + page_size - 1) // page_size

sglang/srt/entrypoints/verl_engine.py

@@ -12,18 +12,17 @@
 # limitations under the License.
 # ==============================================================================
 import os
-from typing import Dict, List, Literal, Optional, Tuple, Union
+from typing import Dict, Iterable, List, Literal, Optional, Tuple, Union
 
 import torch
 import torch.distributed as dist
 from PIL.Image import Image
 from torch.distributed.tensor import DeviceMesh, DTensor
 
+from sglang.srt.entrypoints.engine import Engine
 from sglang.srt.entrypoints.http_server_engine import HttpServerEngineAdapter
 from sglang.srt.model_executor.model_runner import LocalSerializedTensor
 from sglang.srt.patch_torch import monkey_patch_torch_reductions
-from sglang.srt.server import Engine
-from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.utils import MultiprocessingSerializer, broadcast_pyobj
 
 
@@ -125,7 +124,7 @@ class VerlEngine:
 
     def update_weights_from_tensor(
         self,
-        named_tensors: List[Tuple[str, torch.Tensor]],
+        named_tensors: Iterable[Tuple[str, torch.Tensor]],
         load_format: Optional[str] = None,
     ):
         # Most naive implementation, can optimize a lot if it is bottleneck
@@ -154,9 +153,12 @@ class VerlEngine:
                         )
                     ],
                     load_format=load_format,
-                    flush_cache=tensor_index == len(named_tensors) - 1,
+                    flush_cache=False,
                 )
 
+        if self._tp_rank == 0:
+            self._engine.tokenizer_manager.flush_cache()
+
     def release_memory_occupation(self):
         if self._tp_rank == 0:
             self._engine.release_memory_occupation()

sglang/srt/layers/activation.py

@@ -21,13 +21,6 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from sglang.srt.utils import is_cuda_available
-
-_is_cuda = is_cuda_available()
-
-if _is_cuda:
-    from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul
-
 from sglang.srt.custom_op import CustomOp
 from sglang.srt.distributed import (
     divide,
@@ -35,7 +28,12 @@ from sglang.srt.distributed import (
     get_tensor_model_parallel_world_size,
 )
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
-from sglang.srt.utils import set_weight_attrs
+from sglang.srt.utils import is_cuda_available, set_weight_attrs
+
+_is_cuda = is_cuda_available()
+
+if _is_cuda:
+    from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul
 
 logger = logging.getLogger(__name__)