sglang 0.5.2rc2__py3-none-any.whl → 0.5.3__py3-none-any.whl

sglang/srt/disaggregation/decode.py

@@ -21,10 +21,11 @@ Life cycle of a request in the decode server
 from __future__ import annotations
 
 import logging
+import time
 from collections import deque
 from dataclasses import dataclass
 from http import HTTPStatus
-from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union
 
 import torch
 from torch.distributed import ProcessGroup
@@ -45,7 +46,7 @@ from sglang.srt.disaggregation.utils import (
     prepare_abort,
 )
 from sglang.srt.layers.dp_attention import get_attention_tp_size
-from sglang.srt.managers.schedule_batch import FINISH_ABORT, ScheduleBatch
+from sglang.srt.managers.schedule_batch import FINISH_ABORT, RequestStage, ScheduleBatch
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.mem_cache.memory_pool import KVCache, ReqToTokenPool
@@ -218,8 +219,10 @@ class DecodePreallocQueue:
 
         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(
+        kv_manager_class: Type[BaseKVManager] = get_kv_class(
+            self.transfer_backend, KVClassType.MANAGER
+        )
+        kv_manager: BaseKVManager = kv_manager_class(
             kv_args,
             DisaggregationMode.DECODE,
             self.scheduler.server_args,
@@ -248,9 +251,10 @@ class DecodePreallocQueue:
                 mgr=self.kv_manager,
                 bootstrap_addr=f"{req.bootstrap_host}:{req.bootstrap_port}",
                 bootstrap_room=req.bootstrap_room,
-                data_parallel_rank=req.data_parallel_rank,
+                prefill_dp_rank=req.data_parallel_rank,
             )
 
+            req.add_latency(RequestStage.DECODE_PREPARE)
             self.queue.append(
                 DecodeRequest(req=req, kv_receiver=kv_receiver, waiting_for_input=False)
             )
@@ -419,8 +423,13 @@ class DecodePreallocQueue:
                 kv_indices, self.token_to_kv_pool_allocator.page_size
             )
             decode_req.kv_receiver.init(page_indices, decode_req.metadata_buffer_index)
+
             preallocated_reqs.append(decode_req)
             indices_to_remove.add(i)
+            decode_req.req.time_stats.decode_transfer_queue_entry_time = (
+                time.perf_counter()
+            )
+            decode_req.req.add_latency(RequestStage.DECODE_BOOTSTRAP)
 
         self.queue = [
             entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
@@ -514,11 +523,19 @@ class DecodePreallocQueue:
                     dtype=torch.int64,
                     device=self.token_to_kv_pool_allocator.device,
                 ),
+                prefix_lens_cpu=torch.tensor(
+                    [0],
+                    dtype=torch.int64,
+                ),
                 seq_lens=torch.tensor(
                     [num_tokens],
                     dtype=torch.int64,
                     device=self.token_to_kv_pool_allocator.device,
                 ),
+                seq_lens_cpu=torch.tensor(
+                    [num_tokens],
+                    dtype=torch.int64,
+                ),
                 last_loc=torch.tensor(
                     [-1],
                     dtype=torch.int64,
@@ -605,16 +622,23 @@ class DecodeTransferQueue:
                 idx = decode_req.metadata_buffer_index
                 (
                     output_id,
+                    cached_tokens,
                     output_token_logprobs_val,
                     output_token_logprobs_idx,
                     output_top_logprobs_val,
                     output_top_logprobs_idx,
+                    output_topk_p,
+                    output_topk_index,
                     output_hidden_states,
                 ) = self.metadata_buffers.get_buf(idx)
 
                 decode_req.req.output_ids.append(output_id[0].item())
+                decode_req.req.cached_tokens = cached_tokens[0].item()
                 if not self.spec_algorithm.is_none():
+                    decode_req.req.output_topk_p = output_topk_p
+                    decode_req.req.output_topk_index = output_topk_index
                     decode_req.req.hidden_states_tensor = output_hidden_states
+
                 if decode_req.req.return_logprob:
                     decode_req.req.output_token_logprobs_val.append(
                         output_token_logprobs_val[0].item()
@@ -635,10 +659,17 @@ class DecodeTransferQueue:
 
                 if hasattr(decode_req.kv_receiver, "clear"):
                     decode_req.kv_receiver.clear()
+                decode_req.kv_receiver = None
+
+                indices_to_remove.add(i)
+                decode_req.req.time_stats.wait_queue_entry_time = time.perf_counter()
 
                 # special handling for sampling_params.max_new_tokens == 1
                 if decode_req.req.sampling_params.max_new_tokens == 1:
                     # finish immediately
+                    decode_req.req.time_stats.forward_entry_time = (
+                        decode_req.req.time_stats.completion_time
+                    ) = time.perf_counter()
                     decode_req.req.check_finished()
                     self.scheduler.stream_output(
                         [decode_req.req], decode_req.req.return_logprob
@@ -646,8 +677,6 @@ class DecodeTransferQueue:
                     self.tree_cache.cache_finished_req(decode_req.req)
                 else:
                     transferred_reqs.append(decode_req.req)
-
-                indices_to_remove.add(i)
             elif poll in [
                 KVPoll.Bootstrapping,
                 KVPoll.WaitingForInput,
@@ -660,6 +689,7 @@ class DecodeTransferQueue:
         for i in indices_to_remove:
             idx = self.queue[i].metadata_buffer_index
             assert idx != -1
+            self.queue[i].req.add_latency(RequestStage.DECODE_TRANSFERRED)
             self.req_to_metadata_buffer_idx_allocator.free(idx)
 
         self.queue = [
@@ -702,12 +732,15 @@ class SchedulerDisaggregationDecodeMixin:
             elif prepare_mlp_sync_flag:
                 batch, _ = self._prepare_idle_batch_and_run(None)
 
-            if batch is None and (
+            queue_size = (
                 len(self.waiting_queue)
                 + len(self.disagg_decode_transfer_queue.queue)
                 + len(self.disagg_decode_prealloc_queue.queue)
-                == 0
-            ):
+            )
+            if self.server_args.disaggregation_decode_enable_offload_kvcache:
+                queue_size += len(self.decode_offload_manager.ongoing_offload)
+
+            if batch is None and queue_size == 0:
                 self.self_check_during_idle()
 
             self.last_batch = batch
@@ -776,12 +809,15 @@ class SchedulerDisaggregationDecodeMixin:
                 )
                 self.process_batch_result(tmp_batch, tmp_result)
 
-            if batch is None and (
+            queue_size = (
                 len(self.waiting_queue)
                 + len(self.disagg_decode_transfer_queue.queue)
                 + len(self.disagg_decode_prealloc_queue.queue)
-                == 0
-            ):
+            )
+            if self.server_args.disaggregation_decode_enable_offload_kvcache:
+                queue_size += len(self.decode_offload_manager.ongoing_offload)
+
+            if batch is None and queue_size == 0:
                 self.self_check_during_idle()
 
             self.last_batch = batch
@@ -851,6 +887,7 @@ class SchedulerDisaggregationDecodeMixin:
             # we can only add at least `num_not_used_batch` new batch to the running queue
             if i < num_not_used_batch:
                 can_run_list.append(req)
+                req.add_latency(RequestStage.DECODE_WAITING)
                 req.init_next_round_input(self.tree_cache)
             else:
                 waiting_queue.append(req)
@@ -859,6 +896,9 @@ class SchedulerDisaggregationDecodeMixin:
         if len(can_run_list) == 0:
             return None
 
+        for req in can_run_list:
+            req.time_stats.forward_entry_time = time.perf_counter()
+
         # construct a schedule batch with those requests and mark as decode
         new_batch = ScheduleBatch.init_new(
             can_run_list,
@@ -884,9 +924,21 @@ class SchedulerDisaggregationDecodeMixin:
             # if there are still retracted requests, we do not allocate new requests
             return
 
-        req_conns = self.disagg_decode_prealloc_queue.pop_preallocated()
-        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)
+        if not hasattr(self, "polling_count"):
+            self.polling_count = 0
+            self.polling_interval = (
+                self.server_args.disaggregation_decode_polling_interval
+            )
+
+        self.polling_count = (self.polling_count + 1) % self.polling_interval
+
+        if self.polling_count % self.polling_interval == 0:
+            req_conns = self.disagg_decode_prealloc_queue.pop_preallocated()
+            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)
+
+        if self.server_args.disaggregation_decode_enable_offload_kvcache:
+            self.decode_offload_manager.check_offload_progress()

sglang/srt/disaggregation/decode_kvcache_offload_manager.py

@@ -0,0 +1,185 @@
+import logging
+import threading
+import time
+
+import torch
+
+from sglang import ServerArgs
+from sglang.srt.managers.cache_controller import HiCacheController
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.srt.mem_cache.memory_pool import (
+    MHATokenToKVPool,
+    MLATokenToKVPool,
+    ReqToTokenPool,
+)
+from sglang.srt.mem_cache.memory_pool_host import (
+    MHATokenToKVPoolHost,
+    MLATokenToKVPoolHost,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class DecodeKVCacheOffloadManager:
+    """Manage decode-side KV cache offloading lifecycle and operations."""
+
+    def __init__(
+        self,
+        req_to_token_pool: ReqToTokenPool,
+        token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
+        tp_group: torch.distributed.ProcessGroup,
+        tree_cache: BasePrefixCache,
+        server_args: ServerArgs,
+    ) -> None:
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+        self.page_size = server_args.page_size
+        self.server_args = server_args
+        self.request_counter = 0
+        self.tree_cache = tree_cache
+        kv_cache = self.token_to_kv_pool_allocator.get_kvcache()
+        if isinstance(kv_cache, MHATokenToKVPool):
+            self.decode_host_mem_pool = MHATokenToKVPoolHost(
+                kv_cache,
+                server_args.hicache_ratio,
+                server_args.hicache_size,
+                self.page_size,
+                server_args.hicache_mem_layout,
+            )
+        elif isinstance(kv_cache, MLATokenToKVPool):
+            self.decode_host_mem_pool = MLATokenToKVPoolHost(
+                kv_cache,
+                server_args.hicache_ratio,
+                server_args.hicache_size,
+                self.page_size,
+                server_args.hicache_mem_layout,
+            )
+        else:
+            raise ValueError("Unsupported KV cache type for decode offload")
+
+        self.tp_group = tp_group
+        self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group)
+        self.cache_controller = HiCacheController(
+            token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+            mem_pool_host=self.decode_host_mem_pool,
+            page_size=self.page_size,
+            tp_group=tp_group,
+            io_backend=server_args.hicache_io_backend,
+            load_cache_event=threading.Event(),
+            storage_backend=server_args.hicache_storage_backend,
+            model_name=server_args.served_model_name,
+            storage_backend_extra_config=server_args.hicache_storage_backend_extra_config,
+        )
+
+        self.ongoing_offload = {}
+        self.ongoing_backup = {}
+        logger.info("Enable offload kv cache for decode side")
+
+    def offload_kv_cache(self, req) -> bool:
+        """Offload a finished request's KV cache to storage."""
+
+        if self.cache_controller is None or self.decode_host_mem_pool is None:
+            return False
+
+        if req.req_pool_idx == -1:
+            return False
+
+        token_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx]
+        if token_indices.dim() == 0 or token_indices.numel() == 0:
+            logger.debug(
+                f"Request {req.rid} has invalid token_indices: {token_indices}"
+            )
+            return False
+
+        tokens = req.origin_input_ids + req.output_ids
+        aligned_len = (len(tokens) // self.page_size) * self.page_size
+        if aligned_len == 0:
+            return False
+
+        token_indices = token_indices[:aligned_len]
+        tokens = tokens[:aligned_len]
+
+        # Asynchronously offload KV cache from device to host by cache controller
+        self.request_counter += 1
+        ack_id = self.request_counter
+        host_indices = self.cache_controller.write(
+            device_indices=token_indices.long(),
+            node_id=ack_id,
+        )
+        if host_indices is None:
+            logger.error(f"Not enough host memory for request {req.rid}")
+            return False
+
+        self.ongoing_offload[ack_id] = (req, host_indices, tokens, time.time())
+        return True
+
+    def check_offload_progress(self):
+        """Check the progress of offload from device to host and backup from host to storage."""
+        cc = self.cache_controller
+
+        qsizes = torch.tensor(
+            [
+                len(cc.ack_write_queue),
+                cc.ack_backup_queue.qsize(),
+            ],
+            dtype=torch.int,
+        )
+        if self.tp_world_size > 1:
+            torch.distributed.all_reduce(
+                qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group
+            )
+
+        n_write, n_backup = map(int, qsizes.tolist())
+        self._check_offload_progress(n_write)
+        self._check_backup_progress(n_backup)
+
+    def _check_offload_progress(self, finish_count):
+        """Check the progress of offload from device to host."""
+        while finish_count > 0:
+            _, finish_event, ack_list = self.cache_controller.ack_write_queue.pop(0)
+            finish_event.synchronize()
+            for ack_id in ack_list:
+                req, host_indices, tokens, start_time = self.ongoing_offload.pop(ack_id)
+
+                # Release device
+                self.tree_cache.cache_finished_req(req)
+
+                # Trigger async backup from host to storage by cache controller
+                self._trigger_backup(req.rid, host_indices, tokens, start_time)
+            finish_count -= 1
+
+    def _check_backup_progress(self, finish_count):
+        """Check the progress of backup from host to storage."""
+        for _ in range(finish_count):
+            storage_operation = self.cache_controller.ack_backup_queue.get()
+            ack_id = storage_operation.id
+            req_id, host_indices, start_time = self.ongoing_backup.pop(ack_id)
+
+            # Release host memory
+            self.decode_host_mem_pool.free(host_indices)
+
+            logger.debug(
+                f"Finished backup request {req_id}, free host memory, len:{len(host_indices)}, cost time:{time.time() - start_time:.2f} seconds."
+            )
+
+    def _trigger_backup(self, req_id, host_indices, tokens, start_time):
+        """Trigger async backup from host to storage by cache controller."""
+
+        # Generate page hashes and write to storage
+        page_hashes = self._compute_prefix_hash(tokens)
+        ack_id = self.cache_controller.write_storage(
+            host_indices,
+            tokens,
+            hash_value=page_hashes,
+        )
+        self.ongoing_backup[ack_id] = (req_id, host_indices, start_time)
+
+    def _compute_prefix_hash(self, tokens):
+        last_hash = ""
+        page_hashes = []
+        for offset in range(0, len(tokens), self.page_size):
+            page_tokens = tokens[offset : offset + self.page_size]
+            last_hash = self.cache_controller.get_hash_str(page_tokens, last_hash)
+            page_hashes.append(last_hash)
+        return page_hashes

sglang/srt/disaggregation/decode_schedule_batch_mixin.py

@@ -76,6 +76,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
             req_pool_indices, dtype=torch.int64, device=self.device
         )
         self.seq_lens = torch.tensor(seq_lens, dtype=torch.int64, device=self.device)
+        self.seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64)
         self.orig_seq_lens = torch.tensor(
             seq_lens, dtype=torch.int32, device=self.device
         )
@@ -110,7 +111,10 @@ class ScheduleBatchDisaggregationDecodeMixin:
             if req.grammar is not None:
                 # FIXME: this try-except block is for handling unexpected xgrammar issue.
                 try:
-                    req.grammar.accept_token(req.output_ids[-1])
+                    # if it is not None, then the grammar is from a retracted request, and we should not
+                    # accept the token as it's already accepted
+                    if req.grammar.current_token is None:
+                        req.grammar.accept_token(req.output_ids[-1])
                 except ValueError as e:
                     # Grammar accept_token can raise ValueError if the token is not in the grammar.
                     # This can happen if the grammar is not set correctly or the token is invalid.
@@ -122,31 +126,39 @@ class ScheduleBatchDisaggregationDecodeMixin:
                 req.grammar.finished = req.finished()
         self.output_ids = torch.tensor(self.output_ids, device=self.device)
 
-        # Simulate the eagle run. We add mock data to hidden states for the
-        # ease of implementation now meaning the first token will have acc rate
-        # of 0.
-        if not self.spec_algorithm.is_none():
+        # Simulate the eagle run.
+        if self.spec_algorithm.is_eagle():
 
             b = len(self.reqs)
-            topk_p = torch.arange(
-                b * server_args.speculative_eagle_topk,
-                0,
-                -1,
-                device=self.device,
-                dtype=torch.float32,
+            topk = server_args.speculative_eagle_topk
+            topk_p = torch.stack(
+                [
+                    torch.as_tensor(
+                        req.output_topk_p[:topk],
+                        device=self.device,
+                        dtype=torch.float32,
+                    )
+                    for req in self.reqs
+                ],
+                dim=0,
             )
-            topk_p = topk_p.reshape(b, server_args.speculative_eagle_topk)
-            topk_p /= b * server_args.speculative_eagle_topk
-            topk_index = torch.arange(
-                b * server_args.speculative_eagle_topk, device=self.device
+            topk_index = torch.stack(
+                [
+                    torch.as_tensor(
+                        req.output_topk_index[:topk],
+                        device=self.device,
+                        dtype=torch.int64,
+                    )
+                    for req in self.reqs
+                ],
+                dim=0,
             )
-            topk_index = topk_index.reshape(b, server_args.speculative_eagle_topk)
 
             hidden_states_list = [req.hidden_states_tensor for req in self.reqs]
             hidden_states = torch.stack(hidden_states_list, dim=0).to(self.device)
 
             # local import to avoid circular import
-            from sglang.srt.speculative.eagle_utils import EagleDraftInput
+            from sglang.srt.speculative.eagle_info import EagleDraftInput
 
             spec_info = EagleDraftInput(
                 topk_p=topk_p,

sglang/srt/disaggregation/fake/conn.py

@@ -62,7 +62,7 @@ class FakeKVReceiver(BaseKVReceiver):
         mgr: BaseKVManager,
         bootstrap_addr: str,
         bootstrap_room: Optional[int] = None,
-        data_parallel_rank: Optional[int] = None,
+        prefill_dp_rank: Optional[int] = None,
     ):
         self.has_init = False