sglang 0.3.0__py3-none-any.whl → 0.3.1.post1__py3-none-any.whl

sglang/srt/managers/controller_multi.py

@@ -71,12 +71,10 @@ class ControllerMulti:
         self,
         server_args: ServerArgs,
         port_args: PortArgs,
-        model_override_args,
     ):
         # Parse args
         self.server_args = server_args
         self.port_args = port_args
-        self.model_override_args = model_override_args
         self.load_balance_method = LoadBalanceMethod.from_str(
             server_args.load_balance_method
         )
@@ -114,7 +112,6 @@ class ControllerMulti:
                 self.server_args,
                 self.port_args,
                 pipe_controller_writer,
-                self.model_override_args,
                 True,
                 gpu_ids,
                 dp_worker_id,
@@ -189,14 +186,13 @@ def start_controller_process(
     server_args: ServerArgs,
     port_args: PortArgs,
     pipe_writer,
-    model_override_args: dict,
 ):
     """Start a controller process."""
 
     configure_logger(server_args)
 
     try:
-        controller = ControllerMulti(server_args, port_args, model_override_args)
+        controller = ControllerMulti(server_args, port_args)
     except Exception:
         pipe_writer.send(get_exception_traceback())
         raise

sglang/srt/managers/controller_single.py

@@ -40,7 +40,6 @@ class ControllerSingle:
         self,
         server_args: ServerArgs,
         port_args: PortArgs,
-        model_override_args: dict,
         gpu_ids: List[int],
         is_data_parallel_worker: bool,
         dp_worker_id: int,
@@ -76,7 +75,6 @@ class ControllerSingle:
                 tp_rank_range,
                 server_args,
                 port_args.nccl_ports[dp_worker_id],
-                model_override_args,
             )
 
         # Launch tp rank 0
@@ -85,7 +83,6 @@ class ControllerSingle:
             0,
             server_args,
             port_args.nccl_ports[dp_worker_id],
-            model_override_args,
         )
         self.tp_cpu_group = self.tp_server.model_runner.tp_group.cpu_group
 
@@ -126,7 +123,6 @@ def start_controller_process(
     server_args: ServerArgs,
     port_args: PortArgs,
     pipe_writer: multiprocessing.connection.Connection,
-    model_override_args: dict,
     is_data_parallel_worker: bool = False,
     gpu_ids: List[int] = None,
     dp_worker_id: int = None,
@@ -149,7 +145,6 @@ def start_controller_process(
         controller = ControllerSingle(
             server_args,
             port_args,
-            model_override_args,
             gpu_ids,
             is_data_parallel_worker,
             dp_worker_id,

sglang/srt/managers/io_struct.py

@@ -20,7 +20,7 @@ processes (TokenizerManager, DetokenizerManager, Controller).
 
 import copy
 import uuid
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
 from sglang.srt.managers.schedule_batch import BaseFinishReason
@@ -43,6 +43,7 @@ class GenerateReqInput:
     # Whether to return logprobs.
     return_logprob: Optional[Union[List[bool], bool]] = None
     # If return logprobs, the start location in the prompt for returning logprobs.
+    # By default, this value is "-1", which means it will only return logprobs for output tokens.
     logprob_start_len: Optional[Union[List[int], int]] = None
     # If return logprobs, the number of top logprobs to return at each position.
     top_logprobs_num: Optional[Union[List[int], int]] = None
@@ -50,6 +51,13 @@ class GenerateReqInput:
     return_text_in_logprobs: bool = False
     # Whether to stream output.
     stream: bool = False
+    # The modalities of the image data [image, multi-images, video]
+    modalities: Optional[List[str]] = None
+
+    is_single: bool = True
+
+    # LoRA related
+    lora_path: Optional[Union[List[Optional[str]], Optional[str]]] = None
 
     def post_init(self):
         if (self.text is None and self.input_ids is None) or (
@@ -177,6 +185,11 @@ class TokenizedGenerateReqInput:
     top_logprobs_num: int
     # Whether to stream output
     stream: bool
+    # Modalities of the input images
+    modalites: Optional[List[str]] = None
+
+    # LoRA related
+    lora_path: Optional[str] = None  # None means just use the base model
 
 
 @dataclass
@@ -190,6 +203,8 @@ class EmbeddingReqInput:
     # Dummy sampling params for compatibility
     sampling_params: Union[List[Dict], Dict] = None
 
+    is_single: bool = True
+
     def post_init(self):
         if (self.text is None and self.input_ids is None) or (
             self.text is not None and self.input_ids is not None

sglang/srt/managers/policy_scheduler.py

@@ -108,18 +108,25 @@ class PrefillAdder:
     def __init__(
         self,
         tree_cache: BasePrefixCache,
+        running_batch: ScheduleBatch,
+        new_token_ratio: float,
         rem_total_tokens: int,
         rem_input_tokens: int,
         rem_chunk_tokens: Optional[int],
         mixed_with_decode_tokens: int = 0,
     ):
         self.tree_cache = tree_cache
+        self.running_batch = running_batch
+        self.new_token_ratio = new_token_ratio
         self.rem_total_tokens = rem_total_tokens - mixed_with_decode_tokens
+        self.rem_total_tokens_ = self.rem_total_tokens
+        self.total_tokens = rem_total_tokens
         self.rem_input_tokens = rem_input_tokens - mixed_with_decode_tokens
         self.rem_chunk_tokens = rem_chunk_tokens
         if self.rem_chunk_tokens is not None:
             self.rem_chunk_tokens -= mixed_with_decode_tokens
 
+        self.req_states = None
         self.can_run_list = []
         self.new_inflight_req = None
         self.log_hit_tokens = 0
@@ -136,16 +143,20 @@ class PrefillAdder:
             )
         )
 
-    def remove_running_tokens(
-        self, running_batch: ScheduleBatch, new_token_ratio: float
-    ):
+    def remove_running_tokens(self, running_batch: ScheduleBatch):
         self.rem_total_tokens -= sum(
             [
                 min(
                     (r.sampling_params.max_new_tokens - len(r.output_ids)),
                     CLIP_MAX_NEW_TOKENS,
                 )
-                * new_token_ratio
+                * self.new_token_ratio
+                for r in running_batch.reqs
+            ]
+        )
+        self.rem_total_tokens_ -= sum(
+            [
+                r.sampling_params.max_new_tokens - len(r.output_ids)
                 for r in running_batch.reqs
             ]
         )
@@ -154,6 +165,7 @@ class PrefillAdder:
         self, prefix_len: int, extend_input_len: int, max_new_tokens: int
     ):
         self.rem_total_tokens -= extend_input_len + max_new_tokens
+        self.rem_total_tokens_ -= extend_input_len + max_new_tokens
         self.rem_input_tokens -= extend_input_len
         if self.rem_chunk_tokens is not None:
             self.rem_chunk_tokens -= extend_input_len
@@ -161,7 +173,29 @@ class PrefillAdder:
         self.log_hit_tokens += prefix_len
         self.log_input_tokens += extend_input_len
 
+    def add_inflight_req_ignore_eos(self, req: Req):
+        truncated = req.extend_input_len > self.rem_chunk_tokens
+        req.extend_input_len = min(req.extend_input_len, self.rem_chunk_tokens)
+        req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
+        self.can_run_list.append(req)
+
+        self._prefill_one_req(
+            0,
+            req.extend_input_len,
+            (
+                min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS)
+                if not truncated
+                else 0
+            ),
+        )
+
+        # Return if chunked prefill not finished
+        return req if truncated else None
+
     def add_inflight_req(self, req: Req):
+        if req.sampling_params.ignore_eos:
+            return self.add_inflight_req_ignore_eos(req)
+
         truncated = req.extend_input_len > self.rem_chunk_tokens
         req.extend_input_len = min(req.extend_input_len, self.rem_chunk_tokens)
         req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
@@ -190,7 +224,90 @@ class PrefillAdder:
             delta = self.tree_cache.dec_lock_ref(last_node)
             self.rem_total_tokens += delta
 
+    def add_one_req_ignore_eos(self, req: Req):
+        def get_req_state(r):
+            new_token_ratio = (
+                1.0 if r.sampling_params.ignore_eos else self.new_token_ratio
+            )
+            tokens_left = r.sampling_params.max_new_tokens * new_token_ratio - len(
+                r.output_ids
+            )
+            tokens_occupied = len(r.origin_input_ids) + len(r.output_ids)
+
+            if tokens_left > 0:
+                return (tokens_left, tokens_occupied)
+
+            return None
+
+        # Quick Check
+        can_run = False
+        if (
+            req.extend_input_len + req.sampling_params.max_new_tokens
+            <= self.rem_total_tokens
+        ):
+            can_run = True
+
+        if not can_run:
+            if self.req_states is None:
+                self.req_states = []
+                if self.running_batch is not None:
+                    for r in self.running_batch.reqs:
+                        state = get_req_state(r)
+                        if state is not None:
+                            self.req_states.append(state)
+                for r in self.can_run_list:
+                    state = get_req_state(r)
+                    if state is not None:
+                        self.req_states.append(state)
+                state = get_req_state(req)
+                if state is not None:
+                    self.req_states.append(state)
+
+                self.req_states.sort(key=lambda x: x[0])
+            else:
+                state = get_req_state(req)
+                if state is not None:
+                    for i, (tokens_left, tokens_occupied) in enumerate(self.req_states):
+                        if tokens_left >= state[0]:
+                            self.req_states.insert(i, state)
+                            break
+                    else:
+                        self.req_states.append(state)
+
+            tokens_freed = 0
+            for i, (tokens_left, tokens_occupied) in enumerate(self.req_states):
+                decode_steps = (
+                    self.req_states[i + 1][0]
+                    if i + 1 < len(self.req_states)
+                    else tokens_left
+                )
+                bs = len(self.req_states) - i
+                if self.total_tokens + tokens_freed - decode_steps * bs <= 0:
+                    return False
+                tokens_freed += tokens_occupied
+
+        if req.extend_input_len <= self.rem_chunk_tokens:
+            self.can_run_list.append(req)
+            self._prefill_one_req(
+                0,
+                req.extend_input_len,
+                min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS),
+            )
+        else:
+            # Chunked prefill
+            trunc_len = self.rem_chunk_tokens
+            req.extend_input_len = trunc_len
+            req.fill_ids = req.fill_ids[:trunc_len]
+            self.can_run_list.append(req)
+            self.new_inflight_req = req
+            self._prefill_one_req(0, trunc_len, 0)
+
+        return True
+
     def add_one_req(self, req: Req):
+        if req.sampling_params.ignore_eos and self.tree_cache.disable:
+            return self.add_one_req_ignore_eos(req)
+
         total_tokens = req.extend_input_len + min(
             req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS
         )
@@ -233,4 +350,4 @@ class PrefillAdder:
                 self.tree_cache.inc_lock_ref(req.last_node)
                 self._prefill_one_req(prefix_len, trunc_len, 0)
 
-        return True
+        return True and not self.no_remaining_tokens()

sglang/srt/managers/schedule_batch.py

@@ -19,7 +19,7 @@ limitations under the License.
 
 import logging
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
 
@@ -29,20 +29,19 @@ from sglang.srt.constrained.jump_forward import JumpForwardMap
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.mem_cache.chunk_cache import ChunkCache
 from sglang.srt.mem_cache.memory_pool import BaseTokenToKVPool, ReqToTokenPool
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
-
-if TYPE_CHECKING:
-    from sglang.srt.layers.sampler import SampleOutput
-
+from sglang.srt.server_args import ServerArgs
 
 INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5
 
 # Put some global args for easy access
 global_server_args_dict = {
-    "disable_flashinfer": False,
-    "disable_flashinfer_sampling": False,
-    "triton_attention_reduce_in_fp32": False,
-    "enable_mla": False,
+    "attention_backend": ServerArgs.attention_backend,
+    "sampling_backend": ServerArgs.sampling_backend,
+    "triton_attention_reduce_in_fp32": ServerArgs.triton_attention_reduce_in_fp32,
+    "enable_mla": ServerArgs.enable_mla,
+    "torchao_config": ServerArgs.torchao_config,
 }
 
 
@@ -53,8 +52,8 @@ class BaseFinishReason:
     def __init__(self, is_error: bool = False):
         self.is_error = is_error
 
-    def __str__(self):
-        raise NotImplementedError("Subclasses must implement this method")
+    def to_json(self):
+        raise NotImplementedError()
 
 
 class FINISH_MATCHED_TOKEN(BaseFinishReason):
@@ -62,40 +61,57 @@ class FINISH_MATCHED_TOKEN(BaseFinishReason):
         super().__init__()
         self.matched = matched
 
-    def __str__(self) -> str:
-        return f"FINISH_MATCHED_TOKEN: {self.matched}"
+    def to_json(self):
+        return {
+            "type": "stop",  # to match OpenAI API's return value
+            "matched": self.matched,
+        }
 
 
-class FINISH_LENGTH(BaseFinishReason):
-    def __init__(self, length: int):
+class FINISH_MATCHED_STR(BaseFinishReason):
+    def __init__(self, matched: str):
         super().__init__()
-        self.length = length
+        self.matched = matched
 
-    def __str__(self) -> str:
-        return f"FINISH_LENGTH: {self.length}"
+    def to_json(self):
+        return {
+            "type": "stop",  # to match OpenAI API's return value
+            "matched": self.matched,
+        }
 
 
-class FINISH_MATCHED_STR(BaseFinishReason):
-    def __init__(self, matched: str):
+class FINISH_LENGTH(BaseFinishReason):
+    def __init__(self, length: int):
         super().__init__()
-        self.matched = matched
+        self.length = length
 
-    def __str__(self) -> str:
-        return f"FINISH_MATCHED_STR: {self.matched}"
+    def to_json(self):
+        return {
+            "type": "length",  # to match OpenAI API's return value
+            "length": self.length,
+        }
 
 
 class FINISH_ABORT(BaseFinishReason):
     def __init__(self):
         super().__init__(is_error=True)
 
-    def __str__(self) -> str:
-        return "FINISH_ABORT"
+    def to_json(self):
+        return {
+            "type": "abort",
+        }
 
 
 class Req:
     """Store all inforamtion of a request."""
 
-    def __init__(self, rid, origin_input_text, origin_input_ids):
+    def __init__(
+        self,
+        rid: str,
+        origin_input_text: str,
+        origin_input_ids: Tuple[int],
+        lora_path: Optional[str] = None,
+    ):
         # Input and output info
         self.rid = rid
         self.origin_input_text = origin_input_text
@@ -103,10 +119,15 @@ class Req:
         self.origin_input_ids = origin_input_ids
         self.output_ids = []  # Each decode stage's output ids
         self.fill_ids = None  # fill_ids = origin_input_ids + output_ids
+        self.lora_path = lora_path
 
         # Memory info
         self.req_pool_idx = None
 
+        # Check finish
+        self.tokenizer = None
+        self.finished_reason = None
+
         # For incremental decoding
         # ----- | --------- read_ids -------|
         # ----- |   surr_ids  |
@@ -125,38 +146,43 @@ class Req:
         # this does not include the jump forward tokens.
         self.completion_tokens_wo_jump_forward = 0
 
-        # For vision input
+        # For vision inputs
         self.pixel_values = None
         self.image_sizes = None
         self.image_offsets = None
         self.pad_value = None
+        self.modalities = None
 
         # Prefix info
-        self.extend_input_len = 0
         self.prefix_indices = []
+        self.extend_input_len = 0
         self.last_node = None
 
         # Sampling parameters
         self.sampling_params = None
         self.stream = False
 
-        # Check finish
-        self.tokenizer = None
-        self.finished_reason = None
-
-        # Logprobs
+        # Logprobs (arguments)
         self.return_logprob = False
-        self.embedding = None
         self.logprob_start_len = 0
         self.top_logprobs_num = 0
+
+        # Logprobs (return value)
         self.normalized_prompt_logprob = None
         self.input_token_logprobs = None
         self.input_top_logprobs = None
         self.output_token_logprobs = []
         self.output_top_logprobs = []
+
+        # Logprobs (internal values)
         # The tokens is prefilled but need to be considered as decode tokens
         # and should be updated for the decode logprobs
         self.last_update_decode_tokens = 0
+        # The relative logprob_start_len in an extend batch
+        self.extend_logprob_start_len = 0
+
+        # Embedding
+        self.embedding = None
 
         # Constrained decoding
         self.regex_fsm: RegexGuide = None
@@ -333,6 +359,9 @@ class ScheduleBatch:
     token_to_kv_pool: BaseTokenToKVPool
     tree_cache: BasePrefixCache
 
+    forward_mode: ForwardMode = None
+    sampling_info: SamplingBatchInfo = None
+
     # Batched arguments to model runner
     input_ids: torch.Tensor = None
     req_pool_indices: torch.Tensor = None
@@ -343,14 +372,19 @@ class ScheduleBatch:
 
     # For mixed chunekd prefill
     prefix_lens_cpu: List[int] = None
+    running_bs: int = None
 
     # For processing logprobs
     return_logprob: bool = False
     top_logprobs_nums: List[int] = None
 
+    # Stream
+    has_stream: bool = False
+
     @classmethod
     def init_new(cls, reqs, req_to_token_pool, token_to_kv_pool, tree_cache):
         return_logprob = any(req.return_logprob for req in reqs)
+        has_stream = any(req.stream for req in reqs)
 
         return cls(
             reqs=reqs,
@@ -358,18 +392,15 @@ class ScheduleBatch:
             token_to_kv_pool=token_to_kv_pool,
             tree_cache=tree_cache,
             return_logprob=return_logprob,
+            has_stream=has_stream,
         )
 
     def batch_size(self):
-        return len(self.reqs) if self.reqs is not None else 0
+        return len(self.reqs)
 
     def is_empty(self):
         return len(self.reqs) == 0
 
-    def has_stream(self) -> bool:
-        # Return whether batch has at least 1 streaming request
-        return any(r.stream for r in self.reqs)
-
     def alloc_req_slots(self, num_reqs):
         req_pool_indices = self.req_to_token_pool.alloc(num_reqs)
         if req_pool_indices is None:
@@ -396,6 +427,8 @@ class ScheduleBatch:
         return out_cache_loc
 
     def prepare_for_extend(self, vocab_size: int):
+        self.forward_mode = ForwardMode.EXTEND
+
         bs = self.batch_size()
         reqs = self.reqs
         input_ids = [r.fill_ids[len(r.prefix_indices) :] for r in reqs]
@@ -410,8 +443,8 @@ class ScheduleBatch:
         for i, req in enumerate(reqs):
             req.req_pool_idx = req_pool_indices_cpu[i]
             pre_len, seq_len = len(req.prefix_indices), len(req.fill_ids)
-            ext_len = seq_len - pre_len
             seq_lens.append(seq_len)
+            assert seq_len - pre_len == req.extend_input_len
 
             if pre_len > 0:
                 self.req_to_token_pool.req_to_token[req.req_pool_idx][
@@ -419,9 +452,19 @@ class ScheduleBatch:
                 ] = req.prefix_indices
 
             self.req_to_token_pool.req_to_token[req.req_pool_idx][pre_len:seq_len] = (
-                out_cache_loc[pt : pt + ext_len]
+                out_cache_loc[pt : pt + req.extend_input_len]
             )
-            pt += ext_len
+
+            # Compute the relative logprob_start_len in an extend batch
+            if req.logprob_start_len >= pre_len:
+                extend_logprob_start_len = min(
+                    req.logprob_start_len - pre_len, req.extend_input_len - 1
+                )
+            else:
+                extend_logprob_start_len = req.extend_input_len - 1
+
+            req.extend_logprob_start_len = extend_logprob_start_len
+            pt += req.extend_input_len
 
         # Set fields
         with torch.device("cuda"):
@@ -434,18 +477,13 @@ class ScheduleBatch:
         self.out_cache_loc = out_cache_loc
         self.top_logprobs_nums = [r.top_logprobs_num for r in reqs]
         self.prefix_lens_cpu = [len(r.prefix_indices) for r in reqs]
-
+        self.extend_lens_cpu = [r.extend_input_len for r in reqs]
+        self.extend_logprob_start_lens_cpu = [r.extend_logprob_start_len for r in reqs]
         self.sampling_info = SamplingBatchInfo.from_schedule_batch(self, vocab_size)
 
     def mix_with_running(self, running_batch: "ScheduleBatch"):
-        # NOTE: prefix_indices is what has been cached, but we don't cache each decode step
-        prefix_lens_cpu = [len(r.prefix_indices) for r in self.reqs]
-        prefix_lens_cpu.extend(
-            [
-                len(r.origin_input_ids) + len(r.output_ids) - 1
-                for r in running_batch.reqs
-            ]
-        )
+        self.forward_mode = ForwardMode.MIXED
+        running_bs = running_batch.batch_size()
 
         for req in running_batch.reqs:
             req.fill_ids = req.origin_input_ids + req.output_ids
@@ -453,12 +491,22 @@ class ScheduleBatch:
 
         input_ids = torch.cat([self.input_ids, running_batch.input_ids])
         out_cache_loc = torch.cat([self.out_cache_loc, running_batch.out_cache_loc])
-        extend_num_tokens = self.extend_num_tokens + running_batch.batch_size()
+        extend_num_tokens = self.extend_num_tokens + running_bs
+
         self.merge(running_batch)
         self.input_ids = input_ids
         self.out_cache_loc = out_cache_loc
         self.extend_num_tokens = extend_num_tokens
-        self.prefix_lens_cpu = prefix_lens_cpu
+
+        # NOTE: prefix_indices is what has been cached, but we don't cache each decode step
+        self.prefix_lens_cpu.extend(
+            [
+                len(r.origin_input_ids) + len(r.output_ids) - 1
+                for r in running_batch.reqs
+            ]
+        )
+        self.extend_lens_cpu.extend([1] * running_bs)
+        self.extend_logprob_start_lens_cpu.extend([0] * running_bs)
 
     def check_decode_mem(self):
         bs = self.batch_size()
@@ -625,6 +673,8 @@ class ScheduleBatch:
         return jump_forward_reqs
 
     def prepare_for_decode(self, input_ids=None):
+        self.forward_mode = ForwardMode.DECODE
+
         if input_ids is None:
             input_ids = [
                 r.output_ids[-1] if r.output_ids else r.origin_input_ids[-1]
@@ -644,8 +694,6 @@ class ScheduleBatch:
             self.req_pool_indices, self.seq_lens - 1
         ] = self.out_cache_loc
 
-        self.sampling_info.update_regex_vocab_mask(self)
-
     def filter_batch(self, unfinished_indices: List[int]):
         if unfinished_indices is None or len(unfinished_indices) == 0:
             # Filter out all requests
@@ -665,6 +713,7 @@ class ScheduleBatch:
         self.out_cache_loc = 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)
+        self.has_stream = any(req.stream for req in self.reqs)
 
         self.sampling_info.filter(unfinished_indices, new_indices)
 
@@ -675,7 +724,6 @@ class ScheduleBatch:
         self.sampling_info.merge(other.sampling_info)
 
         self.reqs.extend(other.reqs)
-
         self.req_pool_indices = torch.concat(
             [self.req_pool_indices, other.req_pool_indices]
         )
@@ -686,18 +734,4 @@ class ScheduleBatch:
         self.out_cache_loc = None
         self.top_logprobs_nums.extend(other.top_logprobs_nums)
         self.return_logprob = any(req.return_logprob for req in self.reqs)
-
-    def check_sample_results(self, sample_output: SampleOutput):
-        if not torch.all(sample_output.success):
-            probs = sample_output.probs
-            batch_next_token_ids = sample_output.batch_next_token_ids
-            logging.warning("Sampling failed, fallback to top_k=1 strategy")
-            probs = probs.masked_fill(torch.isnan(probs), 0.0)
-            argmax_ids = torch.argmax(probs, dim=-1)
-            batch_next_token_ids = torch.where(
-                sample_output.success, batch_next_token_ids, argmax_ids
-            )
-            sample_output.probs = probs
-            sample_output.batch_next_token_ids = batch_next_token_ids
-
-        return sample_output.batch_next_token_ids
+        self.has_stream = any(req.stream for req in self.reqs)