sglang 0.2.10__py3-none-any.whl → 0.2.12__py3-none-any.whl

sglang/srt/managers/tp_worker.py

@@ -17,35 +17,40 @@ limitations under the License.
 
 import logging
 import multiprocessing
+import os
 import pickle
 import time
 import warnings
-from typing import List, Optional
+from typing import Any, List, Optional, Union
 
 import torch
+import torch.distributed
 import torch.distributed as dist
 
 from sglang.global_config import global_config
 from sglang.srt.constrained.fsm_cache import FSMCache
 from sglang.srt.constrained.jump_forward import JumpForwardCache
 from sglang.srt.hf_transformers_utils import get_processor, get_tokenizer
+from sglang.srt.layers.logits_processor import LogitProcessorOutput
 from sglang.srt.managers.io_struct import (
     AbortReq,
+    BatchEmbeddingOut,
     BatchTokenIDOut,
     FlushCacheReq,
+    TokenizedEmbeddingReqInput,
     TokenizedGenerateReqInput,
 )
-from sglang.srt.managers.policy_scheduler import PolicyScheduler
+from sglang.srt.managers.policy_scheduler import PolicyScheduler, PrefillAdder
 from sglang.srt.managers.schedule_batch import (
     FINISH_ABORT,
     BaseFinishReason,
-    Batch,
-    ForwardMode,
     Req,
+    ScheduleBatch,
 )
 from sglang.srt.mem_cache.chunk_cache import ChunkCache
 from sglang.srt.mem_cache.radix_cache import RadixCache
 from sglang.srt.model_config import ModelConfig
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
@@ -59,6 +64,9 @@ from sglang.utils import get_exception_traceback
 logger = logging.getLogger(__name__)
 
 
+crash_on_warning = os.getenv("SGLANG_IS_IN_CI", "false") == "true"
+
+
 class ModelTpServer:
     def __init__(
         self,
@@ -98,26 +106,24 @@ class ModelTpServer:
             nccl_port=nccl_port,
             server_args=server_args,
         )
-
-        if is_multimodal_model(server_args.model_path):
-            self.processor = get_processor(
-                server_args.tokenizer_path,
-                tokenizer_mode=server_args.tokenizer_mode,
-                trust_remote_code=server_args.trust_remote_code,
-            )
-            self.tokenizer = self.processor.tokenizer
+        if server_args.skip_tokenizer_init:
+            self.tokenizer = self.processor = None
         else:
-            self.tokenizer = get_tokenizer(
-                server_args.tokenizer_path,
-                tokenizer_mode=server_args.tokenizer_mode,
-                trust_remote_code=server_args.trust_remote_code,
-            )
+            if is_multimodal_model(server_args.model_path):
+                self.processor = get_processor(
+                    server_args.tokenizer_path,
+                    tokenizer_mode=server_args.tokenizer_mode,
+                    trust_remote_code=server_args.trust_remote_code,
+                )
+                self.tokenizer = self.processor.tokenizer
+            else:
+                self.tokenizer = get_tokenizer(
+                    server_args.tokenizer_path,
+                    tokenizer_mode=server_args.tokenizer_mode,
+                    trust_remote_code=server_args.trust_remote_code,
+                )
         self.max_total_num_tokens = self.model_runner.max_total_num_tokens
-        self.max_prefill_tokens = (
-            16384
-            if server_args.max_prefill_tokens is None
-            else server_args.max_prefill_tokens
-        )
+        self.max_prefill_tokens = server_args.max_prefill_tokens
         self.max_running_requests = min(
             (
                 self.max_total_num_tokens // 2
@@ -160,19 +166,13 @@ class ModelTpServer:
                 disable=server_args.disable_radix_cache,
             )
         self.tree_cache_metrics = {"total": 0, "hit": 0}
-        self.scheduler = PolicyScheduler(
-            self.schedule_policy,
-            self.max_running_requests,
-            self.max_prefill_tokens,
-            self.max_total_num_tokens,
-            self.tree_cache,
-        )
+        self.scheduler = PolicyScheduler(self.schedule_policy, self.tree_cache)
         self.req_to_token_pool = self.model_runner.req_to_token_pool
         self.token_to_kv_pool = self.model_runner.token_to_kv_pool
 
         # Init running status
         self.waiting_queue: List[Req] = []
-        self.running_batch: Batch = None
+        self.running_batch: ScheduleBatch = None
         self.out_pyobjs = []
         self.decode_forward_ct = 0
         self.stream_interval = server_args.stream_interval
@@ -180,13 +180,15 @@ class ModelTpServer:
         self.last_stats_tic = time.time()
 
         # Init the FSM cache for constrained generation
-        self.regex_fsm_cache = FSMCache(
-            server_args.tokenizer_path,
-            {
-                "tokenizer_mode": server_args.tokenizer_mode,
-                "trust_remote_code": server_args.trust_remote_code,
-            },
-        )
+        if not server_args.skip_tokenizer_init:
+            self.regex_fsm_cache = FSMCache(
+                server_args.tokenizer_path,
+                {
+                    "tokenizer_mode": server_args.tokenizer_mode,
+                    "trust_remote_code": server_args.trust_remote_code,
+                },
+                skip_tokenizer_init=server_args.skip_tokenizer_init,
+            )
         self.jump_forward_cache = JumpForwardCache()
 
         # Init new token estimation
@@ -200,13 +202,14 @@ class ModelTpServer:
         )
         self.new_token_ratio = self.min_new_token_ratio
         self.new_token_ratio_decay = global_config.new_token_ratio_decay
-        self.new_token_ratio_recovery = global_config.new_token_ratio_recovery
 
-    def exposed_step(self, recv_reqs):
+    def exposed_step(self, recv_reqs: List):
         try:
             # Recv requests
             for recv_req in recv_reqs:
-                if isinstance(recv_req, TokenizedGenerateReqInput):
+                if isinstance(
+                    recv_req, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput)
+                ):
                     self.handle_generate_request(recv_req)
                 elif isinstance(recv_req, FlushCacheReq):
                     self.flush_cache()
@@ -233,8 +236,6 @@ class ModelTpServer:
         if new_batch is not None:
             # Run a new prefill batch
             self.forward_prefill_batch(new_batch)
-            self.cache_filled_batch(new_batch)
-            self.filter_out_inflight(new_batch)
 
             if not new_batch.is_empty():
                 if self.running_batch is None:
@@ -251,7 +252,7 @@ class ModelTpServer:
 
                     # Print stats
                     if self.tp_rank == 0 and self.decode_forward_ct % 40 == 0:
-                        self.print_stats()
+                        self.print_decode_stats()
 
                     if self.running_batch.is_empty():
                         self.running_batch = None
@@ -263,7 +264,7 @@ class ModelTpServer:
                 self.check_memory()
                 self.new_token_ratio = global_config.init_new_token_ratio
 
-    def print_stats(self):
+    def print_decode_stats(self):
         num_used = self.max_total_num_tokens - (
             self.token_to_kv_pool.available_size() + self.tree_cache.evictable_size()
         )
@@ -289,52 +290,55 @@ class ModelTpServer:
                 f"available_size={available_size}, max_total_num_tokens={self.max_total_num_tokens}\n"
                 "KV cache pool leak detected!"
             )
+            exit(1) if crash_on_warning else None
 
-        if self.req_to_token_pool.can_use_mem_size != self.req_to_token_pool.size:
+        if len(self.req_to_token_pool.free_slots) != self.req_to_token_pool.size:
             warnings.warn(
                 "Warning: "
-                f"available req slots={self.req_to_token_pool.can_use_mem_size}, "
+                f"available req slots={len(self.req_to_token_pool.free_slots)}, "
                 f"total slots={self.req_to_token_pool.size}\n"
                 "Memory pool leak detected!"
             )
+            exit(1) if crash_on_warning else None
 
     def handle_generate_request(
         self,
-        recv_req: TokenizedGenerateReqInput,
+        recv_req: Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput],
     ):
         req = Req(recv_req.rid, recv_req.input_text, recv_req.input_ids)
-        req.pixel_values = recv_req.pixel_values
-        if req.pixel_values is not None:
-            req.pad_value = [
-                (recv_req.image_hash) % self.model_config.vocab_size,
-                (recv_req.image_hash >> 16) % self.model_config.vocab_size,
-                (recv_req.image_hash >> 32) % self.model_config.vocab_size,
-                (recv_req.image_hash >> 64) % self.model_config.vocab_size,
-            ]
-            req.image_size = recv_req.image_size
-            (
-                req.origin_input_ids,
-                req.image_offset,
-            ) = self.model_runner.model.pad_input_ids(
-                req.origin_input_ids_unpadded,
-                req.pad_value,
-                req.pixel_values.shape,
-                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
-        req.top_logprobs_num = recv_req.top_logprobs_num
-        req.stream = recv_req.stream
         req.tokenizer = self.tokenizer
-
-        # Init regex fsm
-        if req.sampling_params.regex is not None:
-            req.regex_fsm = self.regex_fsm_cache.query(req.sampling_params.regex)
-            if not self.disable_regex_jump_forward:
-                req.jump_forward_map = self.jump_forward_cache.query(
-                    req.sampling_params.regex
+        req.sampling_params = recv_req.sampling_params
+        if self.model_runner.is_generation:
+            req.pixel_values = recv_req.pixel_values
+            if req.pixel_values is not None:
+                req.pad_value = [
+                    (recv_req.image_hash) % self.model_config.vocab_size,
+                    (recv_req.image_hash >> 16) % self.model_config.vocab_size,
+                    (recv_req.image_hash >> 32) % self.model_config.vocab_size,
+                    (recv_req.image_hash >> 64) % self.model_config.vocab_size,
+                ]
+                req.image_size = recv_req.image_size
+                (
+                    req.origin_input_ids,
+                    req.image_offset,
+                ) = self.model_runner.model.pad_input_ids(
+                    req.origin_input_ids_unpadded,
+                    req.pad_value,
+                    req.pixel_values.shape,
+                    req.image_size,
                 )
+            req.return_logprob = recv_req.return_logprob
+            req.logprob_start_len = recv_req.logprob_start_len
+            req.top_logprobs_num = recv_req.top_logprobs_num
+            req.stream = recv_req.stream
+
+            # Init regex fsm
+            if req.sampling_params.regex is not None:
+                req.regex_fsm = self.regex_fsm_cache.query(req.sampling_params.regex)
+                if not self.disable_regex_jump_forward:
+                    req.jump_forward_map = self.jump_forward_cache.query(
+                        req.sampling_params.regex
+                    )
 
         # Truncate prompts that are too long
         if len(req.origin_input_ids) >= self.max_req_input_len:
@@ -343,189 +347,91 @@ class ModelTpServer:
                 "the max context length. Truncated!!!"
             )
             req.origin_input_ids = req.origin_input_ids[: self.max_req_input_len]
-        req.sampling_params.max_new_tokens = min(
-            (
-                req.sampling_params.max_new_tokens
-                if req.sampling_params.max_new_tokens is not None
-                else 1 << 30
-            ),
-            self.max_req_input_len - 1 - len(req.origin_input_ids),
-        )
+
+        if self.model_runner.is_generation:
+            req.sampling_params.max_new_tokens = min(
+                (
+                    req.sampling_params.max_new_tokens
+                    if req.sampling_params.max_new_tokens is not None
+                    else 1 << 30
+                ),
+                self.max_req_input_len - 1 - len(req.origin_input_ids),
+            )
+
         self.waiting_queue.append(req)
 
-    def get_new_prefill_batch(self) -> Optional[Batch]:
-        # TODO(lsyin): organize this function
+    def get_new_prefill_batch(self) -> Optional[ScheduleBatch]:
         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.waiting_queue:
-            req.input_ids = req.origin_input_ids + req.output_ids
-            prefix_indices, last_node = self.tree_cache.match_prefix(
-                rid=req.rid,
-                key=req.input_ids,
-            )
-            if req.return_logprob:
-                prefix_indices = prefix_indices[: req.logprob_start_len]
-            req.extend_input_len = len(req.input_ids) - len(prefix_indices)
-            req.prefix_indices = prefix_indices
-            req.last_node = last_node
+            return None
 
         # Get priority queue
-        self.waiting_queue = self.scheduler.get_priority_queue(self.waiting_queue)
+        prefix_computed = self.scheduler.calc_priority(self.waiting_queue)
 
-        # Add requests if there is available space
-        can_run_list = []
-        new_batch_total_tokens = 0
-        new_batch_input_tokens = 0
-
-        available_size = (
-            self.token_to_kv_pool.available_size() + self.tree_cache.evictable_size()
+        adder = PrefillAdder(
+            self.tree_cache,
+            self.token_to_kv_pool.available_size() + self.tree_cache.evictable_size(),
+            self.max_prefill_tokens,
+            self.chunked_prefill_size,
         )
-        if self.running_batch:
-            available_size -= sum(
-                [
-                    (r.sampling_params.max_new_tokens - len(r.output_ids))
-                    * self.new_token_ratio
-                    for r in self.running_batch.reqs
-                ]
-            )
 
-        # Handle the current inflight request
-        take_inflight = 0
-        if self.current_inflight_req:
-            take_inflight = 1
-            r = self.current_inflight_req
-            r.input_ids = r.origin_input_ids + r.output_ids
-            truncated = (
-                len(r.input_ids) - len(r.prefix_indices) > self.chunked_prefill_size
+        if self.running_batch is not None:
+            adder.remove_running_tokens(self.running_batch, self.new_token_ratio)
+
+        has_inflight = self.current_inflight_req is not None
+        if self.current_inflight_req is not None:
+            self.current_inflight_req.init_next_round_input(
+                None if prefix_computed else self.tree_cache
             )
-            r.extend_input_len = min(
-                len(r.input_ids) - len(r.prefix_indices), self.chunked_prefill_size
+            self.current_inflight_req = adder.add_inflight_req(
+                self.current_inflight_req
             )
-            r.input_ids = r.input_ids[: len(r.prefix_indices) + r.extend_input_len]
-            can_run_list.append(r)
-
-            if not truncated:
-                # Finish inflight
-                self.current_inflight_req = None
-                new_batch_total_tokens += (
-                    r.extend_input_len + r.sampling_params.max_new_tokens
-                )
-                new_batch_input_tokens += r.extend_input_len
-            else:
-                new_batch_total_tokens += r.extend_input_len
-                new_batch_input_tokens += r.extend_input_len
 
         for req in self.waiting_queue:
-            if req.return_logprob and req.normalized_prompt_logprob is None:
-                # Need at least two tokens to compute normalized logprob
-                if req.extend_input_len < 2:
-                    delta = 2 - req.extend_input_len
-                    req.extend_input_len += delta
-                    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.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]
-                if req.image_offset is not None:
-                    req.image_offset += 1
-
+            req.init_next_round_input(None if prefix_computed else self.tree_cache)
+            res = adder.add_one_req(req)
             if (
-                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
-                    <= self.max_prefill_tokens
-                    or len(can_run_list) == 0
-                )
+                not res
+                or adder.no_remaining_tokens()
+                or running_bs + len(adder.can_run_list) >= self.max_running_requests
             ):
-                delta = self.tree_cache.inc_lock_ref(req.last_node)
-                available_size += delta
-
-                if not (
-                    req.extend_input_len
-                    + req.sampling_params.max_new_tokens
-                    + new_batch_total_tokens
-                    < available_size
-                ):
-                    # Undo locking
-                    delta = self.tree_cache.dec_lock_ref(req.last_node)
-                    available_size += delta
-                    break
-                else:
-                    # Add this request to the running batch
-                    if (
-                        self.chunked_prefill_size is None
-                        or (
-                            new_batch_input_tokens + req.extend_input_len
-                            <= self.chunked_prefill_size
-                        )
-                        or (
-                            req.return_logprob and req.normalized_prompt_logprob is None
-                        )
-                    ):
-                        can_run_list.append(req)
-                        new_batch_total_tokens += (
-                            req.extend_input_len + req.sampling_params.max_new_tokens
-                        )
-                        new_batch_input_tokens += req.extend_input_len
-                    else:
-                        trunc_len = self.chunked_prefill_size - new_batch_input_tokens
-
-                        if trunc_len <= 0:
-                            # Undo locking
-                            delta = self.tree_cache.dec_lock_ref(req.last_node)
-                            available_size += delta
-                            break
-
-                        req.extend_input_len = trunc_len
-                        req.input_ids = req.input_ids[
-                            : len(req.prefix_indices) + req.extend_input_len
-                        ]
-                        can_run_list.append(req)
-                        self.current_inflight_req = req
-                        new_batch_input_tokens += req.extend_input_len
-                        new_batch_total_tokens += req.extend_input_len
-                        break
-            else:
                 break
 
-            if running_bs + len(can_run_list) >= self.max_running_requests:
-                break
+        can_run_list = adder.can_run_list
+
+        if adder.new_inflight_req is not None:
+            assert self.current_inflight_req is None
+            self.current_inflight_req = adder.new_inflight_req
 
         if len(can_run_list) == 0:
             return None
 
         # Print stats
         if self.tp_rank == 0:
-            hit_tokens = sum(len(x.prefix_indices) for x in can_run_list)
-            self.tree_cache_metrics["total"] += (
-                hit_tokens + new_batch_input_tokens
-            ) / 10**9
-            self.tree_cache_metrics["hit"] += hit_tokens / 10**9
-            tree_cache_hit_rate = (
-                self.tree_cache_metrics["hit"] / self.tree_cache_metrics["total"]
-            )
+            if isinstance(self.tree_cache, RadixCache):
+                self.tree_cache_metrics["total"] += (
+                    adder.log_input_tokens + adder.log_hit_tokens
+                ) / 10**9
+                self.tree_cache_metrics["hit"] += (adder.log_hit_tokens) / 10**9
+                tree_cache_hit_rate = (
+                    self.tree_cache_metrics["hit"] / self.tree_cache_metrics["total"]
+                )
+            else:
+                tree_cache_hit_rate = 0.0
             logger.info(
                 f"[gpu={self.gpu_id}] Prefill batch. "
                 f"#new-seq: {len(can_run_list)}, "
-                f"#new-token: {new_batch_input_tokens}, "
-                f"#cached-token: {hit_tokens}, "
+                f"#new-token: {adder.log_input_tokens}, "
+                f"#cached-token: {adder.log_hit_tokens}, "
                 f"cache hit rate: {100.0 * tree_cache_hit_rate:.2f}%, "
                 f"#running-req: {running_bs}, "
-                f"#queue-req: {len(self.waiting_queue) - len(can_run_list) + take_inflight}"
+                f"#queue-req: {len(self.waiting_queue) - len(can_run_list) + has_inflight}"
             )
 
         # Return the new batch
-        new_batch = Batch.init_new(
+        new_batch = ScheduleBatch.init_new(
             can_run_list,
             self.req_to_token_pool,
             self.token_to_kv_pool,
@@ -534,47 +440,94 @@ class ModelTpServer:
         self.waiting_queue = [x for x in self.waiting_queue if x not in can_run_list]
         return new_batch
 
-    def forward_prefill_batch(self, batch: Batch):
+    def forward_prefill_batch(self, batch: ScheduleBatch):
         # Build batch tensors
         batch.prepare_for_extend(
             self.model_config.vocab_size, self.int_token_logit_bias
         )
 
-        # Forward and sample the next tokens
-        if batch.extend_num_tokens != 0:
-            output = self.model_runner.forward(batch, ForwardMode.EXTEND)
-            next_token_ids = batch.sample(output.next_token_logits)
-
-            # Move logprobs to cpu
-            if output.next_token_logprobs is not None:
-                output.next_token_logprobs = output.next_token_logprobs[
-                    torch.arange(len(next_token_ids), device=next_token_ids.device),
-                    next_token_ids,
-                ].tolist()
-                output.input_token_logprobs = output.input_token_logprobs.tolist()
-                output.normalized_prompt_logprobs = (
-                    output.normalized_prompt_logprobs.tolist()
-                )
+        if self.model_runner.is_generation:
+            # Forward and sample the next tokens
+            if batch.extend_num_tokens != 0:
+                output = self.model_runner.forward(batch, ForwardMode.EXTEND)
+                next_token_ids = batch.sample(output.next_token_logits)
+
+                # Move logprobs to cpu
+                if output.next_token_logprobs is not None:
+                    output.next_token_logprobs = output.next_token_logprobs[
+                        torch.arange(len(next_token_ids), device=next_token_ids.device),
+                        next_token_ids,
+                    ].tolist()
+                    output.input_token_logprobs = output.input_token_logprobs.tolist()
+                    output.normalized_prompt_logprobs = (
+                        output.normalized_prompt_logprobs.tolist()
+                    )
 
-            next_token_ids = next_token_ids.tolist()
-        else:
-            next_token_ids = [self.tokenizer.eos_token_id] * len(batch.reqs)
+                next_token_ids = next_token_ids.tolist()
+            else:
+                if self.tokenizer is None:
+                    next_token_ids = []
+                    for req in batch.reqs:
+                        next_token_ids.append(
+                            next(iter(req.sampling_params.stop_token_ids))
+                        )
+                else:
+                    next_token_ids = [self.tokenizer.eos_token_id] * len(batch.reqs)
+
+            # Check finish conditions
+            pt = 0
+            for i, req in enumerate(batch.reqs):
+                if req is not self.current_inflight_req:
+                    # Inflight reqs' prefill is not finished
+                    req.completion_tokens_wo_jump_forward += 1
+                    req.output_ids.append(next_token_ids[i])
+                    req.check_finished()
+
+                if req.finished():
+                    self.tree_cache.cache_finished_req(req)
+                else:
+                    self.tree_cache.cache_unfinished_req(req)
 
-        # Check finish conditions
-        pt = 0
-        for i, req in enumerate(batch.reqs):
-            if req is not self.current_inflight_req:
-                req.completion_tokens_wo_jump_forward += 1
-                req.output_ids.append(next_token_ids[i])
-                req.check_finished()
+                if req is self.current_inflight_req:
+                    # Inflight request would get a new req idx
+                    self.req_to_token_pool.free(req.req_pool_idx)
 
-            if req.return_logprob:
-                self.add_logprob_return_values(i, req, pt, next_token_ids, output)
-                pt += req.extend_input_len
+                if req.return_logprob:
+                    self.add_logprob_return_values(i, req, pt, next_token_ids, output)
+                    pt += req.extend_input_len
+        else:
+            assert batch.extend_num_tokens != 0
+            output = self.model_runner.forward(batch, ForwardMode.EXTEND)
+            embeddings = output.embeddings.tolist()
+
+            # Check finish conditions
+            for i, req in enumerate(batch.reqs):
+                req.embedding = embeddings[i]
+                if req is not self.current_inflight_req:
+                    # Inflight reqs' prefill is not finished
+                    # dummy output token for embedding models
+                    req.output_ids.append(0)
+                    req.check_finished()
+
+                if req.finished():
+                    self.tree_cache.cache_finished_req(req)
+                else:
+                    self.tree_cache.cache_unfinished_req(req)
+
+                if req is self.current_inflight_req:
+                    # Inflight request would get a new req idx
+                    self.req_to_token_pool.free(req.req_pool_idx)
 
         self.handle_finished_requests(batch)
 
-    def add_logprob_return_values(self, i, req, pt, next_token_ids, output):
+    def add_logprob_return_values(
+        self,
+        i,
+        req: Req,
+        pt: int,
+        next_token_ids: List[int],
+        output: LogitProcessorOutput,
+    ):
         if req.normalized_prompt_logprob is None:
             req.normalized_prompt_logprob = output.normalized_prompt_logprobs[i]
 
@@ -583,12 +536,12 @@ class ModelTpServer:
             req.input_token_logprobs = list(
                 zip(
                     output.input_token_logprobs[pt : pt + req.extend_input_len - 1],
-                    req.input_ids[-req.extend_input_len + 1 :],
+                    req.fill_ids[-req.extend_input_len + 1 :],
                 )
             )
             if req.logprob_start_len == 0:
                 req.input_token_logprobs = [
-                    (None, req.input_ids[0])
+                    (None, req.fill_ids[0])
                 ] + req.input_token_logprobs
 
         if req.last_update_decode_tokens != 0:
@@ -602,7 +555,7 @@ class ModelTpServer:
                             + req.extend_input_len
                             - 1
                         ],
-                        req.input_ids[-req.last_update_decode_tokens + 1 :],
+                        req.fill_ids[-req.last_update_decode_tokens + 1 :],
                     )
                 )
             )
@@ -623,24 +576,7 @@ class ModelTpServer:
                 )
             req.output_top_logprobs.append(output.output_top_logprobs[i])
 
-    def cache_filled_batch(self, batch: Batch):
-        req_pool_indices_cpu = batch.req_pool_indices.cpu().numpy()
-        for i, req in enumerate(batch.reqs):
-            new_prefix_indices, new_last_node = self.tree_cache.cache_req(
-                rid=req.rid,
-                token_ids=tuple(req.input_ids),
-                last_uncached_pos=len(req.prefix_indices),
-                req_pool_idx=req_pool_indices_cpu[i],
-                del_in_memory_pool=False,
-                old_last_node=req.last_node,
-            )
-            req.prefix_indices, req.last_node = new_prefix_indices, new_last_node
-
-            if req is self.current_inflight_req:
-                # inflight request would get a new req idx
-                self.req_to_token_pool.free(int(req_pool_indices_cpu[i]))
-
-    def forward_decode_batch(self, batch: Batch):
+    def forward_decode_batch(self, batch: ScheduleBatch):
         # Check if decode out of memory
         if not batch.check_decode_mem():
             old_ratio = self.new_token_ratio
@@ -690,6 +626,9 @@ class ModelTpServer:
             req.output_ids.append(next_token_id)
             req.check_finished()
 
+            if req.finished():
+                self.tree_cache.cache_finished_req(req)
+
             if req.return_logprob:
                 req.output_token_logprobs.append(
                     (next_token_logprobs[i], next_token_id)
@@ -699,22 +638,23 @@ class ModelTpServer:
 
         self.handle_finished_requests(batch)
 
-    def handle_finished_requests(self, batch: Batch):
+    def handle_finished_requests(self, batch: ScheduleBatch):
         output_rids = []
-        output_vids = []
-        decoded_texts = []
-        output_read_ids = []
-        output_read_offsets = []
-        output_skip_special_tokens = []
-        output_spaces_between_special_tokens = []
         output_meta_info = []
         output_finished_reason: List[BaseFinishReason] = []
-        finished_indices = []
+        if self.model_runner.is_generation:
+            output_vids = []
+            decoded_texts = []
+            output_read_ids = []
+            output_read_offsets = []
+            output_skip_special_tokens = []
+            output_spaces_between_special_tokens = []
+        else:  # for embedding model
+            output_embeddings = []
         unfinished_indices = []
+
         for i, req in enumerate(batch.reqs):
-            if req.finished():
-                finished_indices.append(i)
-            else:
+            if not req.finished() and req is not self.current_inflight_req:
                 unfinished_indices.append(i)
 
             if req.finished() or (
@@ -727,86 +667,75 @@ class ModelTpServer:
                 )
             ):
                 output_rids.append(req.rid)
-                output_vids.append(req.vid)
-                decoded_texts.append(req.decoded_text)
-                read_ids, read_offset = req.init_incremental_detokenize()
-                output_read_ids.append(read_ids)
-                output_read_offsets.append(read_offset)
-                output_skip_special_tokens.append(
-                    req.sampling_params.skip_special_tokens
-                )
-                output_spaces_between_special_tokens.append(
-                    req.sampling_params.spaces_between_special_tokens
-                )
-
-                meta_info = {
-                    "prompt_tokens": len(req.origin_input_ids),
-                    "completion_tokens": len(req.output_ids),
-                    "completion_tokens_wo_jump_forward": req.completion_tokens_wo_jump_forward,
-                    "finish_reason": str(req.finished_reason),
-                }
-                if req.return_logprob:
-                    (
-                        meta_info["input_token_logprobs"],
-                        meta_info["output_token_logprobs"],
-                        meta_info["input_top_logprobs"],
-                        meta_info["output_top_logprobs"],
-                        meta_info["normalized_prompt_logprob"],
-                    ) = (
-                        req.input_token_logprobs,
-                        req.output_token_logprobs,
-                        req.input_top_logprobs,
-                        req.output_top_logprobs,
-                        req.normalized_prompt_logprob,
-                    )
-                output_meta_info.append(meta_info)
                 output_finished_reason.append(req.finished_reason)
+                if self.model_runner.is_generation:
+                    output_vids.append(req.vid)
+                    decoded_texts.append(req.decoded_text)
+                    read_ids, read_offset = req.init_incremental_detokenize()
+                    output_read_ids.append(read_ids)
+                    output_read_offsets.append(read_offset)
+                    output_skip_special_tokens.append(
+                        req.sampling_params.skip_special_tokens
+                    )
+                    output_spaces_between_special_tokens.append(
+                        req.sampling_params.spaces_between_special_tokens
+                    )
+
+                    meta_info = {
+                        "prompt_tokens": len(req.origin_input_ids),
+                        "completion_tokens": len(req.output_ids),
+                        "completion_tokens_wo_jump_forward": req.completion_tokens_wo_jump_forward,
+                        "finish_reason": str(req.finished_reason),
+                    }
+                    if req.return_logprob:
+                        (
+                            meta_info["input_token_logprobs"],
+                            meta_info["output_token_logprobs"],
+                            meta_info["input_top_logprobs"],
+                            meta_info["output_top_logprobs"],
+                            meta_info["normalized_prompt_logprob"],
+                        ) = (
+                            req.input_token_logprobs,
+                            req.output_token_logprobs,
+                            req.input_top_logprobs,
+                            req.output_top_logprobs,
+                            req.normalized_prompt_logprob,
+                        )
+                    output_meta_info.append(meta_info)
+                else:  # for embedding model
+                    output_embeddings.append(req.embedding)
+                    meta_info = {
+                        "prompt_tokens": len(req.origin_input_ids),
+                    }
+                    output_meta_info.append(meta_info)
 
         # Send to detokenizer
         if output_rids:
-            self.out_pyobjs.append(
-                BatchTokenIDOut(
-                    output_rids,
-                    output_vids,
-                    decoded_texts,
-                    output_read_ids,
-                    output_read_offsets,
-                    output_skip_special_tokens,
-                    output_spaces_between_special_tokens,
-                    output_meta_info,
-                    output_finished_reason,
+            if self.model_runner.is_generation:
+                self.out_pyobjs.append(
+                    BatchTokenIDOut(
+                        output_rids,
+                        output_vids,
+                        decoded_texts,
+                        output_read_ids,
+                        output_read_offsets,
+                        output_skip_special_tokens,
+                        output_spaces_between_special_tokens,
+                        output_meta_info,
+                        output_finished_reason,
+                    )
                 )
-            )
-
-        # Remove finished reqs
-        if finished_indices:
-            # Update radix cache
-            req_pool_indices_cpu = batch.req_pool_indices.tolist()
-            for i in finished_indices:
-                req = batch.reqs[i]
-                self.tree_cache.cache_req(
-                    rid=req.rid,
-                    token_ids=tuple(req.origin_input_ids + req.output_ids)[:-1],
-                    last_uncached_pos=len(req.prefix_indices),
-                    req_pool_idx=req_pool_indices_cpu[i],
+            else:  # for embedding model
+                self.out_pyobjs.append(
+                    BatchEmbeddingOut(
+                        output_rids,
+                        output_embeddings,
+                        output_meta_info,
+                        output_finished_reason,
+                    )
                 )
 
-                self.tree_cache.dec_lock_ref(req.last_node)
-
-            # Update batch tensors
-            if unfinished_indices:
-                batch.filter_batch(unfinished_indices)
-            else:
-                batch.reqs = []
-
-    def filter_out_inflight(self, batch: Batch):
-        # TODO(lsyin): reduce the overhead, make a special version for this
-        if self.current_inflight_req is None:
-            return
-
-        to_remove = batch.reqs.index(self.current_inflight_req)
-        unfinished_indices = [i for i in range(len(batch.reqs)) if i != to_remove]
-
+        # Remove finished reqs: update batch tensors
         batch.filter_batch(unfinished_indices)
 
     def flush_cache(self):
@@ -873,7 +802,11 @@ def run_tp_server(
 
 
 def launch_tp_servers(
-    gpu_ids, tp_rank_range, server_args, nccl_port, model_overide_args
+    gpu_ids: List[int],
+    tp_rank_range: List[int],
+    server_args: ServerArgs,
+    nccl_port: int,
+    model_overide_args: dict,
 ):
     """Launch multiple tensor parallel servers."""
     procs = []
@@ -888,7 +821,9 @@ def launch_tp_servers(
     return procs
 
 
-def broadcast_recv_input(data, rank, dist_group):
+def broadcast_recv_input(
+    data: Any, rank: int, dist_group: torch.distributed.ProcessGroup
+):
     """Broadcast inputs from rank=0 to all other ranks with torch.dist backend."""
 
     if rank == 0: