Improve the control of streaming and improve the first token latency in streaming (#117)

python/sglang/srt/managers/router/model_rpc.py

@@ -17,8 +17,8 @@ from sglang.srt.constrained.fsm_cache import FSMCache
 from sglang.srt.hf_transformers_utils import get_processor, get_tokenizer
 from sglang.srt.managers.io_struct import (
     BatchTokenIDOut,
-    TokenizedGenerateReqInput,
     FlushCacheReq,
+    TokenizedGenerateReqInput,
 )
 from sglang.srt.managers.router.infer_batch import Batch, ForwardMode, Req
 from sglang.srt.managers.router.model_runner import ModelRunner
@@ -194,6 +194,9 @@ class ModelRpcServer(rpyc.Service):
                     if self.running_batch.is_empty():
                         self.running_batch = None
                         break
+
+                    if self.out_pyobjs and self.running_batch.reqs[0].stream:
+                        break
             else:
                 # check the available size
                 available_size = (
@@ -208,8 +211,7 @@ class ModelRpcServer(rpyc.Service):
                     )
 
         if self.running_batch is not None and self.tp_rank == 0:
-            if self.decode_forward_ct >= 20:
-                self.decode_forward_ct = 0
+            if self.decode_forward_ct % 20 == 0:
                 num_used = self.max_total_num_token - (
                     self.token_to_kv_pool.available_size()
                     + self.tree_cache.evictable_size()
@@ -225,11 +227,8 @@ class ModelRpcServer(rpyc.Service):
         self,
         recv_req: TokenizedGenerateReqInput,
     ):
-        req = Req(recv_req.rid)
-        req.input_text = recv_req.input_text
-        req.input_ids = recv_req.input_ids
+        req = Req(recv_req.rid, recv_req.input_text, recv_req.input_ids)
         req.pixel_values = recv_req.pixel_values
-        req.image_size = recv_req.image_size
         if req.pixel_values is not None:
             pad_value = [
                 (recv_req.image_hash) % self.model_config.vocab_size,
@@ -240,6 +239,7 @@ class ModelRpcServer(rpyc.Service):
             req.input_ids, req.image_offset = self.model_runner.model.pad_input_ids(
                 req.input_ids, pad_value, req.pixel_values.shape, req.image_size
             )
+            req.image_size = recv_req.image_size
         req.sampling_params = recv_req.sampling_params
         req.return_logprob = recv_req.return_logprob
         req.logprob_start_len = recv_req.logprob_start_len
@@ -327,9 +327,11 @@ class ModelRpcServer(rpyc.Service):
                     req.extend_input_len + req.max_new_tokens() + new_batch_total_tokens
                     < available_size
                 ):
+                    # Undo the insertion
                     delta = self.tree_cache.dec_ref_counter(req.last_node)
                     available_size += delta
                 else:
+                    # Add this request to the running batch
                     self.token_to_kv_pool.add_refs(req.prefix_indices)
                     can_run_list.append(req)
                     new_batch_total_tokens += (
@@ -421,7 +423,7 @@ class ModelRpcServer(rpyc.Service):
                 return
 
         # Update batch tensors
-        self.decode_forward_ct += 1
+        self.decode_forward_ct = (self.decode_forward_ct + 1) % (1 << 30)
         batch.prepare_for_decode()
 
         # Forward
@@ -454,7 +456,13 @@ class ModelRpcServer(rpyc.Service):
                 unfinished_indices.append(i)
 
             if req.finished or (
-                req.stream and self.decode_forward_ct % self.stream_interval == 0
+                (
+                    req.stream
+                    and (
+                        self.decode_forward_ct % self.stream_interval == 0
+                        or len(req.output_ids) == 1
+                    )
+                )
             ):
                 output_rids.append(req.rid)
                 output_tokens.append(req.output_ids)