sglang 0.5.3rc0__py3-none-any.whl → 0.5.4__py3-none-any.whl

sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -1,16 +1,21 @@
 from __future__ import annotations
 
 import logging
-import threading
 import time
 from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
 import torch
 
 from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.environ import envs
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
-from sglang.srt.managers.io_struct import AbortReq, BatchEmbeddingOut, BatchTokenIDOut
+from sglang.srt.managers.io_struct import (
+    AbortReq,
+    BatchEmbeddingOutput,
+    BatchTokenIDOutput,
+)
 from sglang.srt.managers.schedule_batch import BaseFinishReason, Req, ScheduleBatch
+from sglang.srt.utils.common import ceil_div
 
 if TYPE_CHECKING:
     from sglang.srt.managers.scheduler import (
@@ -35,11 +40,13 @@ class SchedulerOutputProcessorMixin:
         self: Scheduler,
         batch: ScheduleBatch,
         result: Union[GenerationBatchResult, EmbeddingBatchResult],
-        launch_done: Optional[threading.Event] = None,
     ):
         skip_stream_req = None
 
         if self.is_generation:
+            if result.copy_done is not None:
+                result.copy_done.synchronize()
+
             (
                 logits_output,
                 next_token_ids,
@@ -52,22 +59,17 @@ class SchedulerOutputProcessorMixin:
                 result.extend_logprob_start_len_per_req,
             )
 
-            if self.enable_overlap:
-                logits_output, next_token_ids, _ = (
-                    self.tp_worker.resolve_last_batch_result(launch_done)
-                )
-            else:
-                # Move next_token_ids and logprobs to cpu
-                next_token_ids = next_token_ids.tolist()
-                if batch.return_logprob:
-                    if logits_output.next_token_logprobs is not None:
-                        logits_output.next_token_logprobs = (
-                            logits_output.next_token_logprobs.tolist()
-                        )
-                    if logits_output.input_token_logprobs is not None:
-                        logits_output.input_token_logprobs = tuple(
-                            logits_output.input_token_logprobs.tolist()
-                        )
+            # Move next_token_ids and logprobs to cpu
+            next_token_ids = next_token_ids.tolist()
+            if batch.return_logprob:
+                if logits_output.next_token_logprobs is not None:
+                    logits_output.next_token_logprobs = (
+                        logits_output.next_token_logprobs.tolist()
+                    )
+                if logits_output.input_token_logprobs is not None:
+                    logits_output.input_token_logprobs = tuple(
+                        logits_output.input_token_logprobs.tolist()
+                    )
 
             hidden_state_offset = 0
 
@@ -75,15 +77,28 @@ class SchedulerOutputProcessorMixin:
             logprob_pt = 0
 
             for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
-                if req.is_retracted:
+                if self.enable_overlap and req.is_retracted and len(req.output_ids) > 0:
+                    req_idx = batch.req_pool_indices[i]
+                    seq_len = len(req.origin_input_ids) + len(req.output_ids)
+                    pos = batch.req_to_token_pool.req_to_token[req_idx][
+                        seq_len - 1 : seq_len
+                    ]
+                    self.token_to_kv_pool_allocator.free(pos)
                     continue
 
-                if self.is_mixed_chunk and self.enable_overlap and req.finished():
+                if (
+                    self.is_mixed_chunk
+                    and self.enable_overlap
+                    and (req.finished() or req.is_retracted)
+                ):
                     # Free the one delayed token for the mixed decode batch
                     j = len(batch.out_cache_loc) - len(batch.reqs) + i
                     self.token_to_kv_pool_allocator.free(batch.out_cache_loc[j : j + 1])
                     continue
 
+                if req.is_retracted:
+                    continue
+
                 if req.is_chunked <= 0:
                     # req output_ids are set here
                     req.output_ids.append(next_token_id)
@@ -91,7 +106,7 @@ class SchedulerOutputProcessorMixin:
 
                     if req.finished():
                         self.tree_cache.cache_finished_req(req)
-                        req.time_stats.completion_time = time.time()
+                        req.time_stats.completion_time = time.perf_counter()
                     elif not batch.decoding_reqs or req not in batch.decoding_reqs:
                         # This updates radix so others can match
                         self.tree_cache.cache_unfinished_req(req)
@@ -101,7 +116,10 @@ class SchedulerOutputProcessorMixin:
                         assert extend_input_len_per_req is not None
                         extend_logprob_start_len = extend_logprob_start_len_per_req[i]
                         extend_input_len = extend_input_len_per_req[i]
-                        num_input_logprobs = extend_input_len - extend_logprob_start_len
+
+                        num_input_logprobs = self._calculate_num_input_logprobs(
+                            req, extend_input_len, extend_logprob_start_len
+                        )
 
                         if req.return_logprob:
                             self.add_logprob_return_values(
@@ -140,7 +158,7 @@ class SchedulerOutputProcessorMixin:
                             logger.error(
                                 f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
                             )
-                            self.abort_request(AbortReq(req.rid))
+                            self.abort_request(AbortReq(rid=req.rid))
                         req.grammar.finished = req.finished()
                 else:
                     # being chunked reqs' prefill is not finished
@@ -156,8 +174,8 @@ class SchedulerOutputProcessorMixin:
                         extend_input_len = extend_input_len_per_req[i]
                         if extend_logprob_start_len < extend_input_len:
                             # Update input logprobs.
-                            num_input_logprobs = (
-                                extend_input_len - extend_logprob_start_len
+                            num_input_logprobs = self._calculate_num_input_logprobs(
+                                req, extend_input_len, extend_logprob_start_len
                             )
                             if req.return_logprob:
                                 self.add_input_logprob_return_values(
@@ -170,11 +188,22 @@ class SchedulerOutputProcessorMixin:
                                 )
                             logprob_pt += num_input_logprobs
 
-            self.set_next_batch_sampling_info_done(batch)
-
         else:  # embedding or reward model
-            embeddings, bid = result.embeddings, result.bid
-            embeddings = embeddings.tolist()
+            is_sparse = envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set()
+
+            embeddings = result.embeddings
+
+            if is_sparse:
+                batch_ids, token_ids = embeddings.indices()
+                values = embeddings.values()
+
+                embeddings = [{} for _ in range(embeddings.size(0))]
+                for i in range(batch_ids.shape[0]):
+                    embeddings[batch_ids[i].item()][token_ids[i].item()] = values[
+                        i
+                    ].item()
+            else:
+                embeddings = embeddings.tolist()
 
             # Check finish conditions
             for i, req in enumerate(batch.reqs):
@@ -197,29 +226,54 @@ class SchedulerOutputProcessorMixin:
 
         self.stream_output(batch.reqs, batch.return_logprob, skip_stream_req)
 
+    def _resolve_spec_overlap_token_ids(
+        self: Scheduler, result: GenerationBatchResult, batch: ScheduleBatch
+    ) -> List[List[int]]:
+        """Resolve the padding next token ids for speculative decoding with overlap."""
+        assert result.next_token_ids.is_cpu
+        assert result.accept_lens.is_cpu
+        assert result.allocate_lens.is_cpu
+
+        next_token_ids = result.next_token_ids.tolist()
+        accept_lens = result.accept_lens.tolist()
+        result.num_accepted_tokens = sum(accept_lens) - len(batch.reqs)
+
+        predict_tokens = []
+        stride = self.draft_worker.speculative_num_draft_tokens
+        for i, req in enumerate(batch.reqs):
+            predict_tokens.append(
+                next_token_ids[i * stride : i * stride + accept_lens[i]]
+            )
+            req.spec_verify_ct += 1
+
+        return predict_tokens
+
     def process_batch_result_decode(
         self: Scheduler,
         batch: ScheduleBatch,
         result: GenerationBatchResult,
-        launch_done: Optional[threading.Event] = None,
     ):
+        if result.copy_done is not None:
+            result.copy_done.synchronize()
+
         logits_output, next_token_ids, can_run_cuda_graph = (
             result.logits_output,
             result.next_token_ids,
             result.can_run_cuda_graph,
         )
-        self.num_generated_tokens += len(batch.reqs)
 
-        if self.enable_overlap:
-            logits_output, next_token_ids, can_run_cuda_graph = (
-                self.tp_worker.resolve_last_batch_result(launch_done)
-            )
-            next_token_logprobs = logits_output.next_token_logprobs
-        elif batch.spec_algorithm.is_none():
-            # spec decoding handles output logprobs inside verify process.
+        if batch.spec_algorithm.is_none():
             next_token_ids = next_token_ids.tolist()
             if batch.return_logprob:
                 next_token_logprobs = logits_output.next_token_logprobs.tolist()
+        elif batch.is_v2_eagle:
+            next_token_ids = self._resolve_spec_overlap_token_ids(result, batch)
+            allocate_lens_list = result.allocate_lens.tolist()
+            accept_lens_list = result.accept_lens.tolist()
+
+        self.num_generated_tokens += len(batch.reqs)
+        if not batch.spec_algorithm.is_none():
+            self.update_spec_metrics(batch.batch_size(), result.num_accepted_tokens)
 
         self.token_to_kv_pool_allocator.free_group_begin()
 
@@ -227,31 +281,74 @@ class SchedulerOutputProcessorMixin:
         # NOTE: the length of reqs and next_token_ids don't match if it is spec decoding.
         # We should ignore using next_token_ids for spec decoding cases.
         for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
-            if req.is_retracted:
-                continue
+            req: Req
+
+            if self.enable_overlap and (req.finished() or req.is_retracted):
+                indices_to_free = None
+                if batch.spec_algorithm.is_eagle():
+                    from sglang.srt.speculative.eagle_info import EagleDraftInput
+
+                    end_p = allocate_lens_list[i]
+                    start_p = end_p - EagleDraftInput.ALLOC_LEN_PER_DECODE
+                    if self.page_size > 1:
+                        start_p = ceil_div(start_p, self.page_size) * self.page_size
+
+                    indices_to_free = self.req_to_token_pool.req_to_token[
+                        req.req_pool_idx
+                    ][start_p:end_p]
 
-            if self.enable_overlap and req.finished():
-                # Free the one extra delayed token
-                if self.page_size == 1:
-                    self.token_to_kv_pool_allocator.free(batch.out_cache_loc[i : i + 1])
                 else:
-                    # Only free when the extra token is in a new page
-                    if (
-                        len(req.origin_input_ids) + len(req.output_ids) - 1
-                    ) % self.page_size == 0:
-                        self.token_to_kv_pool_allocator.free(
-                            batch.out_cache_loc[i : i + 1]
-                        )
+                    if self.page_size == 1:
+                        # Free the one extra delayed token
+                        indices_to_free = batch.out_cache_loc[i : i + 1]
+                    else:
+                        if (
+                            len(req.origin_input_ids) + len(req.output_ids) - 1
+                        ) % self.page_size == 0:
+                            # Only free when the extra token is in a new page
+                            indices_to_free = batch.out_cache_loc[i : i + 1]
+
+                if indices_to_free is not None:
+                    self.token_to_kv_pool_allocator.free(indices_to_free)
                 continue
 
+            if req.is_retracted:
+                continue
+
+            new_accepted_len = 1
             if batch.spec_algorithm.is_none():
-                # speculative worker will solve the output_ids in speculative decoding
                 req.output_ids.append(next_token_id)
+            elif batch.is_v2_eagle:
+                # Only v2 eagle's output_ids are updated here.
+                req.output_ids.extend(next_token_id)
+                new_accepted_len = len(next_token_id)
+
+            req.check_finished(new_accepted_len)
 
-            req.check_finished()
             if req.finished():
-                self.tree_cache.cache_finished_req(req)
-                req.time_stats.completion_time = time.time()
+                if batch.is_v2_eagle and self.cur_batch.forward_mode.is_extend():
+                    # FIXME(lsyin): fix the messy logic here
+                    # 1) when not overlap (v2 impl), we free the extra tokens in the req
+                    # 2) overlap eagle and the current batch is prefill. This seq will not run extra iteration.
+                    start_p = batch.seq_lens_cpu[i] + accept_lens_list[i]
+                    end_p = allocate_lens_list[i]
+
+                    if self.page_size > 1:
+                        start_p = ceil_div(start_p, self.page_size) * self.page_size
+
+                    indices_to_free = self.req_to_token_pool.req_to_token[
+                        req.req_pool_idx
+                    ][start_p:end_p]
+                    self.token_to_kv_pool_allocator.free(indices_to_free)
+
+                if self.server_args.disaggregation_decode_enable_offload_kvcache:
+                    # Asynchronously offload KV cache; cache_finished_req will be called after Device->Host transfer completes
+                    if not self.decode_offload_manager.offload_kv_cache(req):
+                        self.tree_cache.cache_finished_req(req)
+                else:
+                    self.tree_cache.cache_finished_req(req)
+
+                req.time_stats.completion_time = time.perf_counter()
 
             if req.return_logprob and batch.spec_algorithm.is_none():
                 # speculative worker handles logprob in speculative decoding
@@ -287,10 +384,9 @@ class SchedulerOutputProcessorMixin:
                     logger.error(
                         f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
                     )
-                    self.abort_request(AbortReq(req.rid))
+                    self.abort_request(AbortReq(rid=req.rid))
                 req.grammar.finished = req.finished()
 
-        self.set_next_batch_sampling_info_done(batch)
         self.stream_output(batch.reqs, batch.return_logprob)
         self.token_to_kv_pool_allocator.free_group_end()
 
@@ -301,6 +397,153 @@ class SchedulerOutputProcessorMixin:
         ):
             self.log_decode_stats(can_run_cuda_graph, running_batch=batch)
 
+    def _process_input_token_logprobs(
+        self, req: Req, input_token_logprobs: List
+    ) -> None:
+        """Process input token logprobs values and indices."""
+        is_multi_item_scoring = self._is_multi_item_scoring(req)
+
+        # Process logprob values - handle multi-item scoring vs regular requests
+        if is_multi_item_scoring:
+            # Multi-item scoring: use all logprobs as-is
+            req.input_token_logprobs_val = input_token_logprobs
+        else:
+            # Regular request: add None at start, remove last (sampling token)
+            req.input_token_logprobs_val = [None] + input_token_logprobs[:-1]
+
+        # Process logprob indices based on scoring type
+        if is_multi_item_scoring:
+            # Multi-item scoring: only include delimiter token positions
+            relevant_tokens = req.origin_input_ids[req.logprob_start_len :]
+            input_token_logprobs_idx = [
+                token_id
+                for token_id in relevant_tokens
+                if token_id == self.server_args.multi_item_scoring_delimiter
+            ]
+        else:
+            # Regular request: include all tokens from logprob_start_len onwards
+            input_token_logprobs_idx = req.origin_input_ids[req.logprob_start_len :]
+
+        # Clip padded hash values from image tokens to prevent detokenization errors
+        req.input_token_logprobs_idx = [
+            x if x < self.model_config.vocab_size - 1 else 0
+            for x in input_token_logprobs_idx
+        ]
+
+    def _process_input_top_logprobs(self, req: Req) -> None:
+        """Process input top logprobs."""
+        if req.top_logprobs_num <= 0:
+            return
+
+        is_multi_item_scoring = self._is_multi_item_scoring(req)
+
+        # Initialize arrays - multi-item scoring starts empty, others start with None
+        req.input_top_logprobs_val = [] if is_multi_item_scoring else [None]
+        req.input_top_logprobs_idx = [] if is_multi_item_scoring else [None]
+
+        # Extend arrays with temp values
+        for val, idx in zip(
+            req.temp_input_top_logprobs_val,
+            req.temp_input_top_logprobs_idx,
+            strict=True,
+        ):
+            req.input_top_logprobs_val.extend(val)
+            req.input_top_logprobs_idx.extend(idx)
+
+        # Remove last token (sampling token) for non multi-item scoring requests
+        if not is_multi_item_scoring:
+            req.input_top_logprobs_val.pop()
+            req.input_top_logprobs_idx.pop()
+
+        # Clean up temp storage
+        req.temp_input_top_logprobs_idx = None
+        req.temp_input_top_logprobs_val = None
+
+    def _process_input_token_ids_logprobs(self, req: Req) -> None:
+        """Process input token IDs logprobs."""
+        if req.token_ids_logprob is None:
+            return
+
+        is_multi_item_scoring = self._is_multi_item_scoring(req)
+
+        # Initialize arrays - multi-item scoring starts empty, others start with None
+        req.input_token_ids_logprobs_val = [] if is_multi_item_scoring else [None]
+        req.input_token_ids_logprobs_idx = [] if is_multi_item_scoring else [None]
+
+        # Process temp values - convert tensors to lists and extend arrays
+        for val, idx in zip(
+            req.temp_input_token_ids_logprobs_val,
+            req.temp_input_token_ids_logprobs_idx,
+            strict=True,
+        ):
+            val_list = val.tolist() if isinstance(val, torch.Tensor) else val
+            req.input_token_ids_logprobs_val.extend(
+                val_list if isinstance(val_list, list) else [val_list]
+            )
+            req.input_token_ids_logprobs_idx.extend(idx)
+
+        # Remove last token (sampling token) for non multi-item scoring requests
+        if not is_multi_item_scoring:
+            req.input_token_ids_logprobs_val.pop()
+            req.input_token_ids_logprobs_idx.pop()
+
+        # Clean up temp storage
+        req.temp_input_token_ids_logprobs_idx = None
+        req.temp_input_token_ids_logprobs_val = None
+
+    def _calculate_relevant_tokens_len(self, req: Req) -> int:
+        """Calculate the expected length of logprob arrays based on whether multi-item scoring is enabled.
+
+        For multi-item scoring, only delimiter positions have logprobs.
+        For regular requests, all positions from logprob_start_len onwards have logprobs.
+        """
+        is_multi_item_scoring = self._is_multi_item_scoring(req)
+
+        if is_multi_item_scoring:
+            # Multi-item scoring: count delimiter tokens from logprob_start_len onwards
+            relevant_tokens = req.origin_input_ids[req.logprob_start_len :]
+            return sum(
+                1
+                for token_id in relevant_tokens
+                if token_id == self.server_args.multi_item_scoring_delimiter
+            )
+        else:
+            # Regular request: all tokens from logprob_start_len onwards
+            return len(req.origin_input_ids) - req.logprob_start_len
+
+    def _calculate_num_input_logprobs(
+        self, req: Req, extend_input_len: int, extend_logprob_start_len: int
+    ) -> int:
+        """Calculate the number of input logprobs based on whether multi-item scoring is enabled.
+
+        For multi-item scoring, only delimiter positions have logprobs.
+        For regular requests, all positions in the range have logprobs.
+        """
+        is_multi_item_scoring = self._is_multi_item_scoring(req)
+
+        if is_multi_item_scoring:
+            # Multi-item scoring: count delimiter tokens in the relevant portion
+            relevant_tokens = req.origin_input_ids[
+                extend_logprob_start_len:extend_input_len
+            ]
+            return sum(
+                1
+                for token_id in relevant_tokens
+                if token_id == self.server_args.multi_item_scoring_delimiter
+            )
+        else:
+            # Regular request: all tokens in the range
+            return extend_input_len - extend_logprob_start_len
+
+    def _is_multi_item_scoring(self, req: Req) -> bool:
+        """Check if request uses multi-item scoring.
+
+        Multi-item scoring applies to prefill-only requests when a delimiter
+        token is configured. In this mode, only positions containing the
+        delimiter token receive logprobs.
+        """
+        return req.is_prefill_only and self.server_args.multi_item_scoring_delimiter
+
     def add_input_logprob_return_values(
         self: Scheduler,
         i: int,
@@ -369,63 +612,14 @@ class SchedulerOutputProcessorMixin:
             assert req.input_top_logprobs_val is None
             assert req.input_top_logprobs_idx is None
 
-            # Compute input_token_logprobs_val
-            # Always pad the first one with None.
-            req.input_token_logprobs_val = [None]
-            req.input_token_logprobs_val.extend(input_token_logprobs)
-            # The last input logprob is for sampling, so just pop it out.
-            req.input_token_logprobs_val.pop()
+            # Process all input logprob types using helper functions
+            self._process_input_token_logprobs(req, input_token_logprobs)
+            self._process_input_top_logprobs(req)
 
-            # Compute input_token_logprobs_idx
-            input_token_logprobs_idx = req.origin_input_ids[req.logprob_start_len :]
-            # Clip the padded hash values from image tokens.
-            # Otherwise, it will lead to detokenization errors.
-            input_token_logprobs_idx = [
-                x if x < self.model_config.vocab_size - 1 else 0
-                for x in input_token_logprobs_idx
-            ]
-            req.input_token_logprobs_idx = input_token_logprobs_idx
-
-            if req.top_logprobs_num > 0:
-                req.input_top_logprobs_val = [None]
-                req.input_top_logprobs_idx = [None]
-                assert len(req.temp_input_token_ids_logprobs_val) == len(
-                    req.temp_input_token_ids_logprobs_idx
-                )
-                for val, idx in zip(
-                    req.temp_input_top_logprobs_val,
-                    req.temp_input_top_logprobs_idx,
-                    strict=True,
-                ):
-                    req.input_top_logprobs_val.extend(val)
-                    req.input_top_logprobs_idx.extend(idx)
-
-                # Last token is a sample token.
-                req.input_top_logprobs_val.pop()
-                req.input_top_logprobs_idx.pop()
-                req.temp_input_top_logprobs_idx = None
-                req.temp_input_top_logprobs_val = None
-
-            if req.token_ids_logprob is not None:
-                req.input_token_ids_logprobs_val = [None]
-                req.input_token_ids_logprobs_idx = [None]
-
-                for val, idx in zip(
-                    req.temp_input_token_ids_logprobs_val,
-                    req.temp_input_token_ids_logprobs_idx,
-                    strict=True,
-                ):
-                    req.input_token_ids_logprobs_val.extend(val)
-                    req.input_token_ids_logprobs_idx.extend(idx)
-
-                # Last token is a sample token.
-                req.input_token_ids_logprobs_val.pop()
-                req.input_token_ids_logprobs_idx.pop()
-                req.temp_input_token_ids_logprobs_idx = None
-                req.temp_input_token_ids_logprobs_val = None
+            self._process_input_token_ids_logprobs(req)
 
             if req.return_logprob:
-                relevant_tokens_len = len(req.origin_input_ids) - req.logprob_start_len
+                relevant_tokens_len = self._calculate_relevant_tokens_len(req)
                 assert len(req.input_token_logprobs_val) == relevant_tokens_len
                 assert len(req.input_token_logprobs_idx) == relevant_tokens_len
                 if req.top_logprobs_num > 0:
@@ -517,6 +711,7 @@ class SchedulerOutputProcessorMixin:
         skip_req: Optional[Req] = None,
     ):
         rids = []
+        http_worker_ipcs = []
         finished_reasons: List[BaseFinishReason] = []
 
         decoded_texts = []
@@ -531,6 +726,7 @@ class SchedulerOutputProcessorMixin:
         completion_tokens = []
         cached_tokens = []
         spec_verify_ct = []
+        spec_accepted_tokens = []
         output_hidden_states = None
 
         if return_logprob:
@@ -571,18 +767,26 @@ class SchedulerOutputProcessorMixin:
                     # because of the one additional delayed token. This "continue" prevented the dummy output.
                     continue
                 req.finished_output = True
+                if req.finished_len is None:
+                    req.finished_len = len(req.output_ids)
                 should_output = True
             else:
                 if req.stream:
                     stream_interval = (
                         req.sampling_params.stream_interval or self.stream_interval
                     )
+
+                    # origin stream_interval logic
                     should_output = (
                         len(req.output_ids) % stream_interval == 1
                         if not self.model_config.is_multimodal_gen
                         and stream_interval > 1
                         else len(req.output_ids) % stream_interval == 0
                     )
+
+                    if should_output:
+                        # check_match_stop_str_prefix if  tail_str's suffix match stop_str prefix
+                        should_output &= not req.check_match_stop_str_prefix()
                 else:
                     should_output = (
                         len(req.output_ids) % DEFAULT_FORCE_STREAM_INTERVAL == 0
@@ -596,6 +800,7 @@ class SchedulerOutputProcessorMixin:
                     req.send_output_token_logprobs_offset
                 )
                 rids.append(req.rid)
+                http_worker_ipcs.append(req.http_worker_ipc)
                 finished_reasons.append(
                     req.finished_reason.to_json() if req.finished_reason else None
                 )
@@ -607,21 +812,25 @@ class SchedulerOutputProcessorMixin:
                 else:
                     decode_ids_list.append(decode_ids[req.send_decode_id_offset :])
 
+                # Exclude the tokens after stop condition
+                output_ids_ = req.output_ids_through_stop
+
                 req.send_decode_id_offset = len(decode_ids)
                 read_offsets.append(read_offset)
-                output_ids.append(req.output_ids[send_token_offset:])
-                req.send_token_offset = len(req.output_ids)
+                output_ids.append(output_ids_[send_token_offset:])
+                req.send_token_offset = len(output_ids_)
                 skip_special_tokens.append(req.sampling_params.skip_special_tokens)
                 spaces_between_special_tokens.append(
                     req.sampling_params.spaces_between_special_tokens
                 )
                 no_stop_trim.append(req.sampling_params.no_stop_trim)
                 prompt_tokens.append(len(req.origin_input_ids))
-                completion_tokens.append(len(req.output_ids))
+                completion_tokens.append(len(output_ids_))
                 cached_tokens.append(req.cached_tokens)
 
                 if not self.spec_algorithm.is_none():
                     spec_verify_ct.append(req.spec_verify_ct)
+                    spec_accepted_tokens.append(req.spec_accepted_tokens)
 
                 if return_logprob:
                     if (
@@ -708,9 +917,8 @@ class SchedulerOutputProcessorMixin:
             if self.model_config.is_multimodal_gen:
                 return
 
-            self.send_to_detokenizer.send_pyobj(
-                BatchTokenIDOut(
-                    rids,
+            self.send_to_detokenizer.send_output(
+                BatchTokenIDOutput(
                     finished_reasons,
                     decoded_texts,
                     decode_ids_list,
@@ -723,6 +931,7 @@ class SchedulerOutputProcessorMixin:
                     completion_tokens,
                     cached_tokens,
                     spec_verify_ct,
+                    spec_accepted_tokens,
                     input_token_logprobs_val,
                     input_token_logprobs_idx,
                     output_token_logprobs_val,
@@ -735,7 +944,10 @@ class SchedulerOutputProcessorMixin:
                     input_token_ids_logprobs_idx,
                     output_token_ids_logprobs_val,
                     output_token_ids_logprobs_idx,
-                    output_hidden_states,
+                    output_token_entropy_val=None,
+                    output_hidden_states=output_hidden_states,
+                    rids=rids,
+                    http_worker_ipcs=http_worker_ipcs,
                     placeholder_tokens_idx=None,
                     placeholder_tokens_val=None,
                 )
@@ -743,6 +955,7 @@ class SchedulerOutputProcessorMixin:
 
     def stream_output_embedding(self: Scheduler, reqs: List[Req]):
         rids = []
+        http_worker_ipcs = []
         finished_reasons: List[BaseFinishReason] = []
 
         embeddings = []
@@ -751,17 +964,19 @@ class SchedulerOutputProcessorMixin:
         for req in reqs:
             if req.finished():
                 rids.append(req.rid)
+                http_worker_ipcs.append(req.http_worker_ipc)
                 finished_reasons.append(req.finished_reason.to_json())
                 embeddings.append(req.embedding)
                 prompt_tokens.append(len(req.origin_input_ids))
                 cached_tokens.append(req.cached_tokens)
-        self.send_to_detokenizer.send_pyobj(
-            BatchEmbeddingOut(
-                rids,
+        self.send_to_detokenizer.send_output(
+            BatchEmbeddingOutput(
                 finished_reasons,
                 embeddings,
                 prompt_tokens,
                 cached_tokens,
+                rids=rids,
+                http_worker_ipcs=http_worker_ipcs,
                 placeholder_tokens_idx=None,
                 placeholder_tokens_val=None,
             )