sglang 0.1.14__py3-none-any.whl → 0.1.16__py3-none-any.whl

sglang/srt/managers/router/infer_batch.py

@@ -1,24 +1,36 @@
 from dataclasses import dataclass
-from enum import Enum, auto
+from enum import IntEnum, auto
 from typing import List
 
 import numpy as np
 import torch
+
 from sglang.srt.managers.router.radix_cache import RadixCache
 from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
 
 
-class ForwardMode(Enum):
+class ForwardMode(IntEnum):
     PREFILL = auto()
     EXTEND = auto()
     DECODE = auto()
 
 
-class FinishReason(Enum):
-    LENGTH = auto()
+class FinishReason(IntEnum):
     EOS_TOKEN = auto()
+    LENGTH = auto()
     STOP_STR = auto()
 
+    @staticmethod
+    def to_str(reason):
+        if reason == FinishReason.EOS_TOKEN:
+            return None
+        elif reason == FinishReason.LENGTH:
+            return "length"
+        elif reason == FinishReason.STOP_STR:
+            return "stop"
+        else:
+            return None
+
 
 class Req:
     def __init__(self, rid, input_text, input_ids):
@@ -30,6 +42,7 @@ class Req:
         # Since jump forward may retokenize the prompt with partial outputs,
         # we maintain the original prompt length to report the correct usage.
         self.prompt_tokens = len(input_ids)
+
         # The number of decoded tokens for token usage report. Note that
         # this does not include the jump forward tokens.
         self.completion_tokens_wo_jump_forward = 0
@@ -40,11 +53,11 @@ class Req:
         self.image_offset = 0
         self.pad_value = None
 
+        # Sampling parameters
         self.sampling_params = None
-        self.return_logprob = False
-        self.logprob_start_len = 0
         self.stream = False
 
+        # Check finish
         self.tokenizer = None
         self.finished = False
         self.finish_reason = None
@@ -54,11 +67,17 @@ class Req:
         self.prefix_indices = []
         self.last_node = None
 
-        self.logprob = None
-        self.token_logprob = None
-        self.normalized_logprob = None
-
-        # For constrained decoding
+        # Logprobs
+        self.return_logprob = False
+        self.logprob_start_len = 0
+        self.top_logprobs_num = 0
+        self.normalized_prompt_logprob = None
+        self.prefill_token_logprobs = None
+        self.decode_token_logprobs = None
+        self.prefill_top_logprobs = None
+        self.decode_top_logprobs = None
+
+        # Constrained decoding
         self.regex_fsm = None
         self.regex_fsm_state = 0
         self.jump_forward_map = None
@@ -77,6 +96,9 @@ class Req:
         )
         if first_token.startswith("▁"):
             old_output_str = " " + old_output_str
+        if self.input_text is None:
+            # TODO(lmzheng): This can be wrong. Check with Liangsheng.
+            self.input_text = self.tokenizer.decode(self.input_ids)
         new_input_string = (
             self.input_text
             + self.output_and_jump_forward_str
@@ -159,7 +181,10 @@ class Batch:
     out_cache_loc: torch.Tensor = None
     out_cache_cont_start: torch.Tensor = None
     out_cache_cont_end: torch.Tensor = None
+
+    # for processing logprobs
     return_logprob: bool = False
+    top_logprobs_nums: List[int] = None
 
     # for multimodal
     pixel_values: List[torch.Tensor] = None
@@ -229,12 +254,11 @@ class Batch:
         extend_num_tokens = seq_lens.sum() - prefix_lens.sum()
         out_cache_loc = self.token_to_kv_pool.alloc(extend_num_tokens)
         if out_cache_loc is None:
-            if not self.tree_cache.disable:
-                self.tree_cache.evict(extend_num_tokens, self.token_to_kv_pool.free)
-                out_cache_loc = self.token_to_kv_pool.alloc(extend_num_tokens)
+            self.tree_cache.evict(extend_num_tokens, self.token_to_kv_pool.dec_refs)
+            out_cache_loc = self.token_to_kv_pool.alloc(extend_num_tokens)
 
             if out_cache_loc is None:
-                print("Prefill out of memory. This should nerver happen.")
+                print("Prefill out of memory. This should never happen.")
                 self.tree_cache.pretty_print()
                 exit()
 
@@ -245,10 +269,14 @@ class Batch:
             ] = out_cache_loc[pt : pt + extend_lens[i]]
             pt += extend_lens[i]
 
-        # Handle logit bias
-        logit_bias = torch.zeros((bs, vocab_size), dtype=torch.float32, device=device)
+        # Handle logit bias but only allocate when needed
+        logit_bias = None
         for i in range(bs):
             if reqs[i].sampling_params.dtype == "int":
+                if logit_bias is None:
+                    logit_bias = torch.zeros(
+                        (bs, vocab_size), dtype=torch.float32, device=device
+                    )
                 logit_bias[i] = int_token_logit_bias
 
         # Set fields
@@ -266,6 +294,7 @@ class Batch:
         self.position_ids_offsets = position_ids_offsets
         self.extend_num_tokens = extend_num_tokens
         self.out_cache_loc = out_cache_loc
+        self.top_logprobs_nums = [r.top_logprobs_num for r in reqs]
 
         self.temperatures = torch.tensor(
             [r.sampling_params.temperature for r in reqs],
@@ -295,8 +324,8 @@ class Batch:
         if self.token_to_kv_pool.available_size() >= bs:
             return True
 
-        if not self.tree_cache.disable:
-            self.tree_cache.evict(bs, self.token_to_kv_pool.free)
+        self.tree_cache.evict(bs, self.token_to_kv_pool.dec_refs)
+
         if self.token_to_kv_pool.available_size() >= bs:
             return True
 
@@ -310,27 +339,27 @@ class Batch:
         )
 
         retracted_reqs = []
-        seq_lens_np = self.seq_lens.cpu().numpy()
-        req_pool_indices_np = self.req_pool_indices.cpu().numpy()
+        seq_lens_cpu = self.seq_lens.cpu().numpy()
+        req_pool_indices_cpu = self.req_pool_indices.cpu().numpy()
         while self.token_to_kv_pool.available_size() < len(self.reqs):
             idx = sorted_indices.pop()
             req = self.reqs[idx]
             retracted_reqs.append(req)
 
-            self.tree_cache.dec_ref_counter(req.last_node)
+            # TODO: apply more fine-grained retraction
+            last_uncached_pos = len(req.prefix_indices)
+            token_indices = self.req_to_token_pool.req_to_token[
+                req_pool_indices_cpu[idx]
+            ][last_uncached_pos : seq_lens_cpu[idx]]
+            self.token_to_kv_pool.dec_refs(token_indices)
+
+            self.tree_cache.dec_lock_ref(req.last_node)
             req.prefix_indices = None
             req.last_node = None
             req.extend_input_len = 0
             req.output_ids = []
             req.regex_fsm_state = 0
 
-            # TODO: apply more fine-grained retraction
-
-            token_indices = self.req_to_token_pool.req_to_token[
-                req_pool_indices_np[idx]
-            ][: seq_lens_np[idx]]
-            self.token_to_kv_pool.free(token_indices)
-
         self.filter_batch(sorted_indices)
 
         return retracted_reqs
@@ -349,20 +378,18 @@ class Batch:
                     if len(jump_forward_str) <= 1:
                         continue
 
-                    # insert the old request into tree_cache
-                    token_ids_in_memory = tuple(req.input_ids + req.output_ids)[:-1]
                     if req_pool_indices_cpu is None:
-                        req_pool_indices_cpu = self.req_pool_indices.cpu().tolist()
-                    req_pool_idx = req_pool_indices_cpu[i]
-                    indices = self.req_to_token_pool.req_to_token[
-                        req_pool_idx, : len(token_ids_in_memory)
-                    ]
-                    prefix_len = self.tree_cache.insert(
-                        token_ids_in_memory, indices.clone()
+                        req_pool_indices_cpu = self.req_pool_indices.tolist()
+
+                    # insert the old request into tree_cache
+                    self.tree_cache.cache_req(
+                        token_ids=tuple(req.input_ids + req.output_ids)[:-1],
+                        last_uncached_pos=len(req.prefix_indices),
+                        req_pool_idx=req_pool_indices_cpu[i],
                     )
-                    self.token_to_kv_pool.free(indices[:prefix_len])
-                    self.req_to_token_pool.free(req_pool_idx)
-                    self.tree_cache.dec_ref_counter(req.last_node)
+
+                    # unlock the last node
+                    self.tree_cache.dec_lock_ref(req.last_node)
 
                     # jump-forward
                     req.jump_forward_and_retokenize(jump_forward_str, next_state)
@@ -391,7 +418,7 @@ class Batch:
             self.out_cache_loc = self.token_to_kv_pool.alloc(bs)
 
             if self.out_cache_loc is None:
-                print("Decode out of memory. This should nerver happen.")
+                print("Decode out of memory. This should never happen.")
                 self.tree_cache.pretty_print()
                 exit()
 
@@ -415,6 +442,7 @@ class Batch:
         self.prefix_lens = None
         self.position_ids_offsets = self.position_ids_offsets[new_indices]
         self.out_cache_loc = self.out_cache_cont_start = self.out_cache_cont_end = None
+        self.top_logprobs_nums = [self.top_logprobs_nums[i] for i in unfinished_indices]
         self.return_logprob = any(req.return_logprob for req in self.reqs)
 
         for item in [
@@ -425,9 +453,12 @@ class Batch:
             "presence_penalties",
             "logit_bias",
         ]:
-            setattr(self, item, getattr(self, item)[new_indices])
+            self_val = getattr(self, item, None)
+            # logit_bias can be None
+            if self_val is not None:
+                setattr(self, item, self_val[new_indices])
 
-    def merge(self, other):
+    def merge(self, other: "Batch"):
         self.reqs.extend(other.reqs)
 
         self.req_pool_indices = torch.concat(
@@ -439,6 +470,7 @@ class Batch:
             [self.position_ids_offsets, other.position_ids_offsets]
         )
         self.out_cache_loc = self.out_cache_cont_start = self.out_cache_cont_end = None
+        self.top_logprobs_nums.extend(other.top_logprobs_nums)
         self.return_logprob = any(req.return_logprob for req in self.reqs)
 
         for item in [
@@ -447,17 +479,34 @@ class Batch:
             "top_ks",
             "frequency_penalties",
             "presence_penalties",
-            "logit_bias",
         ]:
-            setattr(
-                self, item, torch.concat([getattr(self, item), getattr(other, item)])
+            self_val = getattr(self, item, None)
+            other_val = getattr(other, item, None)
+            setattr(self, item, torch.concat([self_val, other_val]))
+
+        # logit_bias can be None
+        if self.logit_bias is not None or other.logit_bias is not None:
+            vocab_size = (
+                self.logit_bias.shape[1]
+                if self.logit_bias is not None
+                else other.logit_bias.shape[1]
             )
+            if self.logit_bias is None:
+                self.logit_bias = torch.zeros(
+                    (len(self.reqs), vocab_size), dtype=torch.float32, device="cuda"
+                )
+            if other.logit_bias is None:
+                other.logit_bias = torch.zeros(
+                    (len(other.reqs), vocab_size), dtype=torch.float32, device="cuda"
+                )
+            self.logit_bias = torch.concat([self.logit_bias, other.logit_bias])
 
     def sample(self, logits: torch.Tensor):
         # Post process logits
         logits = logits.contiguous()
         logits.div_(self.temperatures)
-        logits.add_(self.logit_bias)
+        if self.logit_bias is not None:
+            logits.add_(self.logit_bias)
 
         has_regex = any(req.regex_fsm is not None for req in self.reqs)
         if has_regex:

sglang/srt/managers/router/manager.py

@@ -4,7 +4,8 @@ import logging
 import uvloop
 import zmq
 import zmq.asyncio
-from sglang.srt.backend_config import GLOBAL_BACKEND_CONFIG
+
+from sglang.global_config import global_config
 from sglang.srt.managers.router.model_rpc import ModelRpcClient
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.utils import get_exception_traceback
@@ -29,7 +30,7 @@ class RouterManager:
         self.recv_reqs = []
 
         # Init some configs
-        self.extend_dependency_time = GLOBAL_BACKEND_CONFIG.extend_dependency_time
+        self.request_dependency_time = global_config.request_dependency_time
 
     async def loop_for_forward(self):
         while True:
@@ -41,12 +42,16 @@ class RouterManager:
                 self.send_to_detokenizer.send_pyobj(obj)
 
             # async sleep for receiving the subsequent request and avoiding cache miss
+            slept = False
             if len(out_pyobjs) != 0:
                 has_finished = any([obj.finished for obj in out_pyobjs])
                 if has_finished:
-                    await asyncio.sleep(self.extend_dependency_time)
+                    if self.request_dependency_time > 0:
+                        slept = True
+                        await asyncio.sleep(self.request_dependency_time)
 
-            await asyncio.sleep(0.0006)
+            if not slept:
+                await asyncio.sleep(0.0006)
 
     async def loop_for_recv_requests(self):
         while True:
@@ -55,9 +60,7 @@ class RouterManager:
 
 
 def start_router_process(
-    server_args: ServerArgs,
-    port_args: PortArgs,
-    pipe_writer,
+    server_args: ServerArgs, port_args: PortArgs, pipe_writer, model_overide_args
 ):
     logging.basicConfig(
         level=getattr(logging, server_args.log_level.upper()),
@@ -65,7 +68,7 @@ def start_router_process(
     )
 
     try:
-        model_client = ModelRpcClient(server_args, port_args)
+        model_client = ModelRpcClient(server_args, port_args, model_overide_args)
         router = RouterManager(model_client, port_args)
     except Exception:
         pipe_writer.send(get_exception_traceback())