sglang 0.3.3__py3-none-any.whl → 0.3.4__py3-none-any.whl

sglang/srt/managers/tokenizer_manager.py

@@ -46,6 +46,7 @@ from sglang.srt.managers.io_struct import (
     EmbeddingReqInput,
     FlushCacheReq,
     GenerateReqInput,
+    ProfileReq,
     RewardReqInput,
     TokenizedEmbeddingReqInput,
     TokenizedGenerateReqInput,
@@ -149,9 +150,13 @@ class TokenizerManager:
         while self.model_update_lock.locked():
             await asyncio.sleep(0.001)
 
+        if isinstance(obj, EmbeddingReqInput) and self.is_generation:
+            raise ValueError(
+                "This model does not appear to be an embedding model by default. Please add `--is-embedding` when launching the server or try another model."
+            )
+
         obj.post_init()
         is_single = obj.is_single
-
         if is_single:
             async for response in self._handle_single_request(obj, request):
                 yield response
@@ -512,6 +517,14 @@ class TokenizerManager:
         req = AbortReq(rid)
         self.send_to_scheduler.send_pyobj(req)
 
+    def start_profile(self):
+        req = ProfileReq.START_PROFILE
+        self.send_to_scheduler.send_pyobj(req)
+
+    def stop_profile(self):
+        req = ProfileReq.STOP_PROFILE
+        self.send_to_scheduler.send_pyobj(req)
+
     async def update_weights(
         self, obj: UpdateWeightReqInput, request: Optional[fastapi.Request] = None
     ):

sglang/srt/managers/tp_worker.py

@@ -17,6 +17,11 @@ limitations under the License.
 
 import json
 import logging
+import threading
+import time
+from queue import Queue
+
+import torch
 
 from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.hf_transformers_utils import get_processor, get_tokenizer
@@ -75,6 +80,7 @@ class TpModelWorker:
                     tokenizer_mode=server_args.tokenizer_mode,
                     trust_remote_code=server_args.trust_remote_code,
                 )
+        self.device = self.model_runner.device
 
         # Profile number of tokens
         self.max_total_num_tokens = self.model_runner.max_total_num_tokens
@@ -100,6 +106,9 @@ class TpModelWorker:
         )[0]
         set_random_seed(self.random_seed)
 
+        if server_args.enable_overlap_schedule:
+            self.init_overlap_status()
+
     def get_token_and_memory_info(self):
         return (
             self.max_total_num_tokens,
@@ -109,6 +118,81 @@ class TpModelWorker:
             self.random_seed,
         )
 
+    def init_overlap_status(self):
+        self.future_logits_output_dict = dict()
+        self.future_logits_output_ct = 0
+        self.future_token_ids_ct = 0
+        self.future_token_ids_map = torch.empty(
+            (self.max_running_requests * 5,), dtype=torch.int32, device=self.device
+        )
+        self.future_token_ids_limit = self.max_running_requests * 3
+        self.future_token_ids_output = dict()
+
+        self.future_event_map = dict()
+        self.forward_queue = Queue()
+        self.forward_stream = torch.cuda.Stream()
+        self.forward_thread = threading.Thread(
+            target=self.forward_thread_func,
+        )
+        self.forward_thread.start()
+
+    def forward_thread_func(self):
+        with torch.cuda.stream(self.forward_stream):
+            self.forward_thread_func_()
+
+    @torch.inference_mode()
+    def forward_thread_func_(self):
+        while True:
+            tic1 = time.time()
+            model_worker_batch, future_logits_output, future_next_token_ids = (
+                self.forward_queue.get()
+            )
+
+            # Resolve future tokens in the input
+            tic2 = time.time()
+            resolved_input_ids = model_worker_batch.input_ids
+            future_mask = resolved_input_ids < 0
+            resolved_input_ids[future_mask] = self.future_token_ids_map[
+                -resolved_input_ids[future_mask]
+            ]
+
+            # Run forward
+            logits_output, next_token_ids = self.forward_batch_generation(
+                model_worker_batch
+            )
+
+            # Set future values
+            if model_worker_batch.return_logprob:
+                self.future_logits_output_dict[future_logits_output] = logits_output
+
+            # logger.info(f"set output {future_next_token_ids=}, {next_token_ids=}")
+            self.future_token_ids_map[-future_next_token_ids] = next_token_ids.to(
+                torch.int32
+            )
+            # logger.info("Set event")
+            self.future_token_ids_output[model_worker_batch.bid] = (
+                next_token_ids.tolist()
+            )
+            self.future_event_map[model_worker_batch.bid].set()
+
+            if False:
+                tic3 = time.time()
+                self.acc_time_with_waiting += tic3 - tic1
+                self.acc_time_without_waiting += tic3 - tic2
+                if self.forward_queue.qsize() == 0:
+                    logger.info(
+                        f"{self.acc_time_with_waiting=:.3f}, {self.acc_time_without_waiting=:.3f}, {self.forward_queue.qsize()=}"
+                    )
+
+    def resolve_future_token_ids(self, bid: int):
+        self.future_event_map[bid].wait()
+        ret = self.future_token_ids_output[bid]
+        del self.future_event_map[bid]
+        return ret
+
+    def resolve_future_logits_output(self, future_obj):
+        return self.future_logits_output_dict.pop(future_obj)
+
     def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch):
         forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
         logits_output = self.model_runner.forward(forward_batch)
@@ -118,9 +202,35 @@ class TpModelWorker:
     def forward_batch_embedding(self, model_worker_batch: ModelWorkerBatch):
         forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
         logits_output = self.model_runner.forward(forward_batch)
-        embeddings = logits_output.embeddings.tolist()
+        embeddings = logits_output.embeddings
         return embeddings
 
+    def forward_batch_generation_non_blocking(
+        self, model_worker_batch: ModelWorkerBatch
+    ):
+        # Allocate output future objects
+        future_logits_output = self.future_logits_output_ct
+        self.future_logits_output_ct += 1
+
+        bs = len(model_worker_batch.seq_lens)
+        with torch.cuda.stream(self.forward_stream):
+            future_next_token_ids = -torch.arange(
+                self.future_token_ids_ct + 1,
+                self.future_token_ids_ct + 1 + bs,
+                dtype=torch.int32,
+                device=self.device,
+            )
+        self.future_token_ids_ct = (
+            self.future_token_ids_ct + bs
+        ) % self.future_token_ids_limit
+        ret = future_logits_output, future_next_token_ids
+
+        self.future_event_map[model_worker_batch.bid] = threading.Event()
+        self.forward_queue.put(
+            (model_worker_batch.copy(), future_logits_output, future_next_token_ids)
+        )
+        return ret
+
     def update_weights(self, recv_req: UpdateWeightReqInput):
         success, message = self.model_runner.update_weights(
             recv_req.model_path, recv_req.load_format

sglang/srt/mem_cache/chunk_cache.py

@@ -40,10 +40,12 @@ class ChunkCache(BasePrefixCache):
 
     def cache_finished_req(self, req: Req, token_ids: Optional[List[int]] = None):
         if token_ids is None:
-            token_ids = (req.origin_input_ids + req.output_ids)[:-1]
+            token_id_len = len(req.origin_input_ids) + len(req.output_ids) - 1
+        else:
+            token_id_len = len(token_ids)
 
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, : len(token_ids)
+            req.req_pool_idx, :token_id_len
         ]
         self.req_to_token_pool.free(req.req_pool_idx)
         self.token_to_kv_pool.free(kv_indices)
@@ -53,10 +55,12 @@ class ChunkCache(BasePrefixCache):
 
     def cache_unfinished_req(self, req: Req, token_ids: Optional[List[int]] = None):
         if token_ids is None:
-            token_ids = req.fill_ids
+            token_id_len = len(req.fill_ids)
+        else:
+            token_id_len = len(token_ids)
 
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, : len(token_ids)
+            req.req_pool_idx, :token_id_len
         ]
 
         if req.rid not in self.entries:

sglang/srt/mem_cache/memory_pool.py

@@ -18,7 +18,6 @@ limitations under the License.
 import logging
 from typing import List, Tuple, Union
 
-import numpy as np
 import torch
 
 logger = logging.getLogger(__name__)
@@ -77,6 +76,8 @@ class BaseTokenToKVPool:
             self.store_dtype = dtype
 
         self.free_slots = None
+        self.is_not_in_free_group = True
+        self.free_group = []
         self.clear()
 
     def available_size(self):
@@ -89,14 +90,28 @@ class BaseTokenToKVPool:
         select_index = self.free_slots[:need_size]
         self.free_slots = self.free_slots[need_size:]
 
-        return torch.tensor(select_index, dtype=torch.int32, device=self.device)
+        return select_index.to(self.device, non_blocking=True)
 
     def free(self, free_index: torch.Tensor):
-        self.free_slots = np.concatenate((self.free_slots, free_index.cpu().numpy()))
+        if self.is_not_in_free_group:
+            self.free_slots = torch.concat((self.free_slots, free_index.cpu()))
+        else:
+            self.free_group.append(free_index)
+
+    def free_group_begin(self):
+        self.is_not_in_free_group = False
+        self.free_group = []
+
+    def free_group_end(self):
+        self.is_not_in_free_group = True
+        if self.free_group:
+            self.free(torch.concat(self.free_group))
 
     def clear(self):
         # The padded slot 0 is used for writing dummy outputs from padded tokens.
-        self.free_slots = np.arange(1, self.size + 1)
+        self.free_slots = torch.arange(1, self.size + 1, dtype=torch.int32)
+        self.is_in_free_group = False
+        self.free_group = []
 
     def get_key_buffer(self, layer_id: int) -> torch.Tensor:
         raise NotImplementedError()
@@ -231,3 +246,61 @@ class MLATokenToKVPool(BaseTokenToKVPool):
             self.kv_buffer[layer_id][loc] = cache_k.view(self.store_dtype)
         else:
             self.kv_buffer[layer_id][loc] = cache_k
+
+
+class DoubleSparseTokenToKVPool(BaseTokenToKVPool):
+
+    def __init__(
+        self,
+        size: int,
+        dtype: torch.dtype,
+        head_num: int,
+        head_dim: int,
+        layer_num: int,
+        device: str,
+        heavy_channel_num: int,
+    ):
+        super().__init__(size, dtype, device)
+
+        # [size, head_num, head_dim] for each layer
+        self.k_buffer = [
+            torch.empty((size + 1, head_num, head_dim), dtype=dtype, device=device)
+            for _ in range(layer_num)
+        ]
+        self.v_buffer = [
+            torch.empty((size + 1, head_num, head_dim), dtype=dtype, device=device)
+            for _ in range(layer_num)
+        ]
+
+        # [size, head_num, heavy_channel_num] for each layer
+        self.label_buffer = [
+            torch.empty(
+                (size + 1, head_num, heavy_channel_num), dtype=dtype, device=device
+            )
+            for _ in range(layer_num)
+        ]
+
+    def get_key_buffer(self, layer_id: int):
+        return self.k_buffer[layer_id]
+
+    def get_value_buffer(self, layer_id: int):
+        return self.v_buffer[layer_id]
+
+    def get_label_buffer(self, layer_id: int):
+        return self.label_buffer[layer_id]
+
+    def get_kv_buffer(self, layer_id: int):
+        return self.k_buffer[layer_id], self.v_buffer[layer_id]
+
+    def set_kv_buffer(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+        cache_label: torch.Tensor,
+    ):
+        # NOTE(Andy): ignore the dtype check
+        self.k_buffer[layer_id][loc] = cache_k
+        self.v_buffer[layer_id][loc] = cache_v
+        self.label_buffer[layer_id][loc] = cache_label

sglang/srt/mem_cache/radix_cache.py

@@ -99,17 +99,25 @@ class RadixCache(BasePrefixCache):
 
     def cache_finished_req(self, req: Req, token_ids: Optional[List[int]] = None):
         """Cache request when it finishes."""
+        if self.disable:
+            if token_ids is None:
+                token_ids_len = len(req.origin_input_ids) + len(req.output_ids) - 1
+            else:
+                token_ids_len = len(token_ids)
+
+            kv_indices = self.req_to_token_pool.req_to_token[
+                req.req_pool_idx, :token_ids_len
+            ]
+            self.token_to_kv_pool.free(kv_indices)
+            self.req_to_token_pool.free(req.req_pool_idx)
+            return
+
         if token_ids is None:
             token_ids = (req.origin_input_ids + req.output_ids)[:-1]
         kv_indices = self.req_to_token_pool.req_to_token[
             req.req_pool_idx, : len(token_ids)
         ]
 
-        if self.disable:
-            self.token_to_kv_pool.free(kv_indices)
-            self.req_to_token_pool.free(req.req_pool_idx)
-            return
-
         # Radix Cache takes one ref in memory pool
         new_prefix_len = self.insert(token_ids, kv_indices.clone())
         self.token_to_kv_pool.free(kv_indices[len(req.prefix_indices) : new_prefix_len])
@@ -229,7 +237,7 @@ class RadixCache(BasePrefixCache):
     def _split_node(self, key, child: TreeNode, split_len: int):
         # new_node -> child
         new_node = TreeNode()
-        new_node.children = {key[split_len:][0]: child}
+        new_node.children = {key[split_len]: child}
         new_node.parent = child.parent
         new_node.lock_ref = child.lock_ref
         new_node.key = child.key[:split_len]
@@ -237,7 +245,7 @@ class RadixCache(BasePrefixCache):
         child.parent = new_node
         child.key = child.key[split_len:]
         child.value = child.value[split_len:]
-        new_node.parent.children[key[:split_len][0]] = new_node
+        new_node.parent.children[key[0]] = new_node
         return new_node
 
     def _insert_helper(self, node: TreeNode, key: List, value):

sglang/srt/model_executor/cuda_graph_runner.py

@@ -245,10 +245,10 @@ class CudaGraphRunner:
             self.out_cache_loc.zero_()
 
         # Common inputs
-        self.input_ids[:raw_bs] = forward_batch.input_ids
-        self.req_pool_indices[:raw_bs] = forward_batch.req_pool_indices
-        self.seq_lens[:raw_bs] = forward_batch.seq_lens
-        self.out_cache_loc[:raw_bs] = forward_batch.out_cache_loc
+        self.input_ids[:raw_bs].copy_(forward_batch.input_ids)
+        self.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices)
+        self.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
+        self.out_cache_loc[:raw_bs].copy_(forward_batch.out_cache_loc)
 
         # Attention backend
         self.model_runner.attn_backend.init_forward_metadata_replay_cuda_graph(

sglang/srt/model_executor/forward_batch_info.py

@@ -118,7 +118,7 @@ class ForwardBatch:
         batch: ModelWorkerBatch,
         model_runner: ModelRunner,
     ):
-        device = "cuda"
+        device = model_runner.device
 
         ret = cls(
             forward_mode=batch.forward_mode,
@@ -134,27 +134,23 @@ class ForwardBatch:
         )
 
         # Init position information
-        if ret.forward_mode.is_decode():
-            ret.positions = (ret.seq_lens - 1).to(torch.int64)
-        else:
-            ret.positions = torch.tensor(
-                np.concatenate(
-                    [
-                        np.arange(prefix_len, prefix_len + extend_len)
-                        for prefix_len, extend_len in zip(
-                            batch.extend_prefix_lens, batch.extend_seq_lens
-                        )
-                    ],
-                    axis=0,
-                ),
-                device=device,
-            ).to(torch.int64)
-
+        if not ret.forward_mode.is_decode():
+            ret.positions = torch.concat(
+                [
+                    torch.arange(prefix_len, prefix_len + extend_len, device=device)
+                    for prefix_len, extend_len in zip(
+                        batch.extend_prefix_lens, batch.extend_seq_lens
+                    )
+                ],
+                axis=0,
+            )
             ret.image_inputs = batch.image_inputs
-            ret.extend_seq_lens = torch.tensor(batch.extend_seq_lens, device=device)
+            ret.extend_seq_lens = torch.tensor(
+                batch.extend_seq_lens, dtype=torch.int32
+            ).to(device, non_blocking=True)
             ret.extend_prefix_lens = torch.tensor(
-                batch.extend_prefix_lens, device=device
-            )
+                batch.extend_prefix_lens, dtype=torch.int32
+            ).to(device, non_blocking=True)
             ret.extend_start_loc = torch.zeros_like(ret.extend_seq_lens)
             ret.extend_start_loc[1:] = torch.cumsum(ret.extend_seq_lens[:-1], dim=0)
             ret.extend_seq_lens_cpu = batch.extend_seq_lens
@@ -164,7 +160,6 @@ class ForwardBatch:
         ret.req_to_token_pool = model_runner.req_to_token_pool
         ret.token_to_kv_pool = model_runner.token_to_kv_pool
         ret.attn_backend = model_runner.attn_backend
-        model_runner.attn_backend.init_forward_metadata(ret)
 
         # Init lora information
         if model_runner.server_args.lora_paths is not None: