sglang 0.1.19__py3-none-any.whl → 0.1.21__py3-none-any.whl

sglang/srt/managers/controller/tp_worker.py

@@ -53,7 +53,7 @@ class ModelTpServer:
         tp_rank: int,
         server_args: ServerArgs,
         model_port_args: ModelPortArgs,
-        model_overide_args,
+        model_overide_args: dict,
     ):
         server_args, model_port_args = obtain(server_args), obtain(model_port_args)
         suppress_other_loggers()
@@ -98,7 +98,7 @@ class ModelTpServer:
             )
         self.max_total_num_tokens = self.model_runner.max_total_num_tokens
         self.max_prefill_tokens = (
-            4096
+            16384
             if server_args.max_prefill_tokens is None
             else server_args.max_prefill_tokens
         )
@@ -178,7 +178,7 @@ class ModelTpServer:
         self.new_token_ratio_recovery = global_config.new_token_ratio_recovery
 
     def exposed_step(self, recv_reqs):
-        if self.tp_size * self.dp_size != 1:
+        if not isinstance(recv_reqs, list):
             recv_reqs = obtain(recv_reqs)
 
         try:
@@ -206,11 +206,11 @@ class ModelTpServer:
 
     @torch.inference_mode()
     def forward_step(self):
-        new_batch = self.get_new_fill_batch()
+        new_batch = self.get_new_prefill_batch()
 
         if new_batch is not None:
-            # Run a new fill batch
-            self.forward_fill_batch(new_batch)
+            # Run a new prefill batch
+            self.forward_prefill_batch(new_batch)
             self.cache_filled_batch(new_batch)
 
             if not new_batch.is_empty():
@@ -219,33 +219,32 @@ class ModelTpServer:
                 else:
                     self.running_batch.merge(new_batch)
         else:
-            # Run decode batch
+            # Run a decode batch
             if self.running_batch is not None:
                 # Run a few decode batches continuously for reducing overhead
-                for _ in range(10):
+                for _ in range(global_config.num_continue_decode_steps):
                     self.num_generated_tokens += len(self.running_batch.reqs)
                     self.forward_decode_batch(self.running_batch)
 
                     # Print stats
-                    if self.tp_rank == 0:
-                        if self.decode_forward_ct % 40 == 0:
-                            num_used = self.max_total_num_tokens - (
-                                self.token_to_kv_pool.available_size()
-                                + self.tree_cache.evictable_size()
-                            )
-                            throughput = self.num_generated_tokens / (
-                                time.time() - self.last_stats_tic
-                            )
-                            self.num_generated_tokens = 0
-                            self.last_stats_tic = time.time()
-                            logger.info(
-                                f"[gpu_id={self.gpu_id}] Decode batch. "
-                                f"#running-req: {len(self.running_batch.reqs)}, "
-                                f"#token: {num_used}, "
-                                f"token usage: {num_used / self.max_total_num_tokens:.2f}, "
-                                f"gen throughput (token/s): {throughput:.2f}, "
-                                f"#queue-req: {len(self.forward_queue)}"
-                            )
+                    if self.tp_rank == 0 and self.decode_forward_ct % 40 == 0:
+                        num_used = self.max_total_num_tokens - (
+                            self.token_to_kv_pool.available_size()
+                            + self.tree_cache.evictable_size()
+                        )
+                        throughput = self.num_generated_tokens / (
+                            time.time() - self.last_stats_tic
+                        )
+                        self.num_generated_tokens = 0
+                        self.last_stats_tic = time.time()
+                        logger.info(
+                            f"[gpu_id={self.gpu_id}] Decode batch. "
+                            f"#running-req: {len(self.running_batch.reqs)}, "
+                            f"#token: {num_used}, "
+                            f"token usage: {num_used / self.max_total_num_tokens:.2f}, "
+                            f"gen throughput (token/s): {throughput:.2f}, "
+                            f"#queue-req: {len(self.forward_queue)}"
+                        )
 
                     if self.running_batch.is_empty():
                         self.running_batch = None
@@ -313,12 +312,12 @@ class ModelTpServer:
         )
         self.forward_queue.append(req)
 
-    def get_new_fill_batch(self) -> Optional[Batch]:
-        if (
-            self.running_batch is not None
-            and len(self.running_batch.reqs) > self.max_running_requests
-        ):
-            return None
+    def get_new_prefill_batch(self) -> Optional[Batch]:
+        running_bs = (
+            len(self.running_batch.reqs) if self.running_batch is not None else 0
+        )
+        if running_bs >= self.max_running_requests:
+            return
 
         # Compute matched prefix length
         for req in self.forward_queue:
@@ -344,7 +343,7 @@ class ModelTpServer:
         if self.running_batch:
             available_size -= sum(
                 [
-                    (r.max_new_tokens() - len(r.output_ids)) * self.new_token_ratio
+                    (r.sampling_params.max_new_tokens - len(r.output_ids)) * self.new_token_ratio
                     for r in self.running_batch.reqs
                 ]
             )
@@ -358,7 +357,7 @@ class ModelTpServer:
                     req.prefix_indices = req.prefix_indices[:-delta]
                     if req.image_offset is not None:
                         req.image_offset += delta
-            if req.extend_input_len == 0 and req.max_new_tokens() > 0:
+            if req.extend_input_len == 0 and req.sampling_params.max_new_tokens > 0:
                 # Need at least one token to compute logits
                 req.extend_input_len = 1
                 req.prefix_indices = req.prefix_indices[:-1]
@@ -366,7 +365,7 @@ class ModelTpServer:
                     req.image_offset += 1
 
             if (
-                req.extend_input_len + req.max_new_tokens() + new_batch_total_tokens
+                req.extend_input_len + req.sampling_params.max_new_tokens + new_batch_total_tokens
                 < available_size
                 and (
                     req.extend_input_len + new_batch_input_tokens
@@ -378,7 +377,7 @@ class ModelTpServer:
                 available_size += delta
 
                 if not (
-                    req.extend_input_len + req.max_new_tokens() + new_batch_total_tokens
+                    req.extend_input_len + req.sampling_params.max_new_tokens + new_batch_total_tokens
                     < available_size
                 ):
                     # Undo locking
@@ -389,19 +388,20 @@ class ModelTpServer:
                     # Add this request to the running batch
                     can_run_list.append(req)
                     new_batch_total_tokens += (
-                        req.extend_input_len + req.max_new_tokens()
+                        req.extend_input_len + req.sampling_params.max_new_tokens
                     )
                     new_batch_input_tokens += req.extend_input_len
             else:
                 break
+
+            if running_bs + len(can_run_list) >= self.max_running_requests:
+                break
+
         if len(can_run_list) == 0:
             return None
 
         # Print stats
         if self.tp_rank == 0:
-            running_req = (
-                0 if self.running_batch is None else len(self.running_batch.reqs)
-            )
             hit_tokens = sum(len(x.prefix_indices) for x in can_run_list)
             self.tree_cache_metrics["total"] += (
                 hit_tokens + new_batch_input_tokens
@@ -416,7 +416,7 @@ class ModelTpServer:
                 f"#new-token: {new_batch_input_tokens}, "
                 f"#cached-token: {hit_tokens}, "
                 f"cache hit rate: {100.0 * tree_cache_hit_rate:.2f}%, "
-                f"#running-req: {running_req}, "
+                f"#running-req: {running_bs}, "
                 f"#queue-req: {len(self.forward_queue) - len(can_run_list)}"
             )
             # logger.debug(
@@ -436,7 +436,7 @@ class ModelTpServer:
         self.forward_queue = [x for x in self.forward_queue if x not in can_run_list]
         return new_batch
 
-    def forward_fill_batch(self, batch: Batch):
+    def forward_prefill_batch(self, batch: Batch):
         # Build batch tensors
         batch.prepare_for_extend(
             self.model_config.vocab_size, self.int_token_logit_bias
@@ -746,8 +746,8 @@ class ModelTpClient:
             # Init model
             assert len(gpu_ids) == 1
             self.model_server = ModelTpService().exposed_ModelTpServer(
-                0,
                 gpu_ids[0],
+                0,
                 server_args,
                 model_port_args,
                 model_overide_args,

sglang/srt/memory_pool.py

@@ -8,96 +8,98 @@ logger = logging.getLogger(__name__)
 
 
 class ReqToTokenPool:
-    def __init__(self, size, max_context_len):
+    """A memory pool that maps a request to its token locations."""
+
+    def __init__(self, size: int, max_context_len: int):
         self.mem_state = torch.ones((size,), dtype=torch.bool, device="cuda")
-        self.can_use_mem_size = size
         self.req_to_token = torch.empty(
             (size, max_context_len), dtype=torch.int32, device="cuda"
         )
+        self.can_use_mem_size = size
 
-    def alloc(self, need_size):
+    def alloc(self, need_size: int):
         if need_size > self.can_use_mem_size:
             return None
 
-        select_index = torch.nonzero(self.mem_state).squeeze(1)[:need_size]
-        self.mem_state[select_index] = 0
+        select_index = torch.nonzero(self.mem_state).squeeze(1)[:need_size].to(torch.int32)
+        self.mem_state[select_index] = False
         self.can_use_mem_size -= need_size
-        return select_index.to(torch.int32)
 
-    def free(self, free_index):
+        return select_index
+
+    def free(self, free_index: int):
+        self.mem_state[free_index] = True
         if isinstance(free_index, (int,)):
             self.can_use_mem_size += 1
         else:
             self.can_use_mem_size += free_index.shape[0]
-        self.mem_state[free_index] = 1
 
     def clear(self):
-        self.mem_state.fill_(1)
+        self.mem_state.fill_(True)
         self.can_use_mem_size = len(self.mem_state)
 
 
 class TokenToKVPool:
+    """A memory pool that maps a token to its kv cache locations"""
+
     def __init__(self, size, dtype, head_num, head_dim, layer_num):
-        self.mem_state = torch.zeros((size,), dtype=torch.int16, device="cuda")
-        self.total_ref_ct = 0
+        self.size = size
+
+        # We also add one slot. This slot is used for writing dummy output from padded tokens.
+        self.mem_state = torch.ones((self.size + 1,), dtype=torch.bool, device="cuda")
 
         # [size, key/value, head_num, head_dim] for each layer
         self.kv_data = [
-            torch.empty((size, 2, head_num, head_dim), dtype=dtype, device="cuda")
+            torch.empty((size + 1, 2, head_num, head_dim), dtype=dtype, device="cuda")
             for _ in range(layer_num)
         ]
 
+        # Prefetch buffer
+        self.prefetch_buffer = torch.empty(0, device="cuda", dtype=torch.int32)
+        self.prefetch_chunk_size = 512
+
+        self.can_use_mem_size = self.size
+        self.clear()
+
     def get_key_buffer(self, layer_id):
         return self.kv_data[layer_id][:, 0]
 
     def get_value_buffer(self, layer_id):
         return self.kv_data[layer_id][:, 1]
 
+    def available_size(self):
+        return self.can_use_mem_size + len(self.prefetch_buffer)
+
     def alloc(self, need_size):
-        select_index = torch.nonzero(self.mem_state == 0).squeeze(1)[:need_size]
-        if select_index.shape[0] < need_size:
-            return None
+        buffer_len = len(self.prefetch_buffer)
+        if need_size <= buffer_len:
+            select_index = self.prefetch_buffer[:need_size]
+            self.prefetch_buffer = self.prefetch_buffer[need_size:]
+            return select_index
 
-        self.add_refs(select_index)
-        return select_index.to(torch.int32)
+        addition_size = need_size - buffer_len
+        alloc_size = max(addition_size, self.prefetch_chunk_size)
+        select_index = torch.nonzero(self.mem_state).squeeze(1)[:alloc_size].to(torch.int32)
 
-    def alloc_contiguous(self, need_size):
-        empty_index = torch.nonzero(self.mem_state == 0).squeeze(1)[:need_size]
-        if empty_index.shape[0] < need_size:
-            return None
-        empty_size = len(empty_index)
-        loc_sum = (
-            empty_index[need_size - 1 :] - empty_index[: empty_size - (need_size - 1)]
-        )
-        can_used_loc = empty_index[: empty_size - (need_size - 1)][
-            loc_sum == need_size - 1
-        ]
-        if can_used_loc.shape[0] == 0:
+        if select_index.shape[0] < addition_size:
             return None
 
-        start_loc = can_used_loc[0].item()
-        select_index = torch.arange(start_loc, start_loc + need_size, device="cuda")
-        self.add_refs(select_index)
-        return select_index.to(torch.int32), start_loc, start_loc + need_size
+        self.mem_state[select_index] = False
+        self.can_use_mem_size -= len(select_index)
 
-    def used_size(self):
-        return len(torch.nonzero(self.mem_state).squeeze(1))
+        self.prefetch_buffer = torch.cat((self.prefetch_buffer, select_index))
+        ret_index = self.prefetch_buffer[:need_size]
+        self.prefetch_buffer = self.prefetch_buffer[need_size:]
 
-    def available_size(self):
-        return torch.sum(self.mem_state == 0).item()
-
-    def add_refs(self, token_index: torch.Tensor):
-        self.total_ref_ct += len(token_index)
-        self.mem_state[token_index] += 1
-
-    def dec_refs(self, token_index: torch.Tensor):
-        self.total_ref_ct -= len(token_index)
-        self.mem_state[token_index] -= 1
+        return ret_index
 
-        num_freed = torch.sum(self.mem_state[token_index] == 0)
-
-        return num_freed
+    def free(self, free_index: torch.Tensor):
+        self.mem_state[free_index] = True
+        self.can_use_mem_size += len(free_index)
 
     def clear(self):
-        self.mem_state.fill_(0)
-        self.total_ref_ct = 0
+        self.mem_state.fill_(True)
+        self.can_use_mem_size = self.size
+
+        # We also add one slot. This slot is used for writing dummy output from padded tokens.
+        self.mem_state[0] = False

sglang/srt/models/minicpm.py

@@ -5,12 +5,9 @@ from typing import Any, Dict, Iterable, Optional, Tuple
 
 import torch
 from torch import nn
-
 from vllm.config import CacheConfig
 from vllm.distributed import get_tensor_model_parallel_world_size
-
 from vllm.model_executor.layers.activation import SiluAndMul
-
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -31,7 +28,6 @@ from sglang.srt.managers.controller.model_runner import InputMetadata
 
 
 class MiniCPMMLP(nn.Module):
-
     def __init__(
         self,
         hidden_size: int,
@@ -67,7 +63,6 @@ class MiniCPMMLP(nn.Module):
 
 
 class MiniCPMAttention(nn.Module):
-
     def __init__(
         self,
         hidden_size: int,
@@ -152,7 +147,6 @@ class MiniCPMAttention(nn.Module):
 
 
 class MiniCPMDecoderLayer(nn.Module):
-
     def __init__(
         self,
         config,
@@ -217,7 +211,6 @@ class MiniCPMDecoderLayer(nn.Module):
 
 
 class MiniCPMModel(nn.Module):
-
     def __init__(
         self,
         config,
@@ -274,7 +267,7 @@ class MiniCPMForCausalLM(nn.Module):
     ) -> None:
         super().__init__()
         self.config = config
-        
+
         self.num_experts = getattr(self.config, "num_experts", 0)
         self.quant_config = quant_config
         self.model = MiniCPMModel(config, quant_config=quant_config)