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

sglang/srt/layers/logits_processor.py

@@ -29,7 +29,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardMode, InputMetad
 
 
 @dataclasses.dataclass
-class LogitProcessorOutput:
+class LogitsProcessorOutput:
     # The logits of the next tokens.       shape: [#seq, vocab_size]
     next_token_logits: torch.Tensor
     # The logprobs of the next tokens.     shape: [#seq, vocab_size]
@@ -55,6 +55,9 @@ class LogitsMetadata:
     extend_start_loc: Optional[torch.Tensor] = None
     top_logprobs_nums: Optional[List[int]] = None
 
+    extend_seq_lens_cpu: List[int] = None
+    logprob_start_lens_cpu: List[int] = None
+
     @classmethod
     def from_input_metadata(cls, input_metadata: InputMetadata):
         return cls(
@@ -63,22 +66,30 @@ class LogitsMetadata:
             extend_start_loc=input_metadata.extend_start_loc,
             return_logprob=input_metadata.return_logprob,
             top_logprobs_nums=input_metadata.top_logprobs_nums,
+            extend_seq_lens_cpu=input_metadata.extend_seq_lens_cpu,
+            logprob_start_lens_cpu=input_metadata.logprob_start_lens_cpu,
         )
 
 
 class LogitsProcessor(nn.Module):
-    def __init__(self, config):
+    def __init__(self, config, skip_all_gather: bool = False):
         super().__init__()
         self.config = config
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.do_tensor_parallel_all_gather = (
+            not skip_all_gather and get_tensor_model_parallel_world_size() > 1
+        )
 
     def _get_normalized_prompt_logprobs(
-        self, input_token_logprobs, logits_metadata: LogitsMetadata
+        self,
+        input_token_logprobs: torch.Tensor,
+        cum_start_len0: torch.Tensor,
+        cum_start_len1: torch.Tensor,
+        logits_metadata: LogitsMetadata,
     ):
         logprobs_cumsum = torch.cumsum(input_token_logprobs, dim=0, dtype=torch.float32)
 
-        start = logits_metadata.extend_start_loc.clone()
-        end = start + logits_metadata.extend_seq_lens - 2
+        start = logits_metadata.extend_start_loc.clone() - cum_start_len0
+        end = start + logits_metadata.extend_seq_lens - 2 - cum_start_len1
         start.clamp_(min=0, max=input_token_logprobs.shape[0] - 1)
         end.clamp_(min=0, max=input_token_logprobs.shape[0] - 1)
         sum_logp = (
@@ -91,7 +102,7 @@ class LogitsProcessor(nn.Module):
         return normalized_prompt_logprobs
 
     @staticmethod
-    def get_top_logprobs(all_logprobs, logits_metadata: LogitsMetadata):
+    def get_top_logprobs(all_logprobs: torch.Tensor, logits_metadata: LogitsMetadata):
         if logits_metadata.forward_mode == ForwardMode.DECODE:
             output_top_logprobs = []
             max_k = max(logits_metadata.top_logprobs_nums)
@@ -105,7 +116,7 @@ class LogitsProcessor(nn.Module):
             # TODO: vectorize the code below
             input_top_logprobs, output_top_logprobs = [], []
             pt = 0
-            extend_seq_lens_cpu = logits_metadata.extend_seq_lens.tolist()
+            extend_seq_lens_cpu = logits_metadata.extend_seq_lens_cpu
 
             max_k = max(logits_metadata.top_logprobs_nums)
             ret = all_logprobs.topk(max_k, dim=1)
@@ -113,26 +124,30 @@ class LogitsProcessor(nn.Module):
             indices = ret.indices.tolist()
 
             for i, extend_seq_len in enumerate(extend_seq_lens_cpu):
+                start_len = logits_metadata.logprob_start_lens_cpu[i]
+                pruned_len = extend_seq_len - start_len
+
                 if extend_seq_len == 0:
                     input_top_logprobs.append([])
                     output_top_logprobs.append([])
                     continue
+
                 k = logits_metadata.top_logprobs_nums[i]
                 input_top_logprobs.append(
                     [
                         list(zip(values[pt + j][:k], indices[pt + j][:k]))
-                        for j in range(extend_seq_len - 1)
+                        for j in range(pruned_len - 1)
                     ]
                 )
                 output_top_logprobs.append(
                     list(
                         zip(
-                            values[pt + extend_seq_len - 1][:k],
-                            indices[pt + extend_seq_len - 1][:k],
+                            values[pt + pruned_len - 1][:k],
+                            indices[pt + pruned_len - 1][:k],
                         )
                     )
                 )
-                pt += extend_seq_len
+                pt += pruned_len
 
             return input_top_logprobs, output_top_logprobs
 
@@ -159,18 +174,18 @@ class LogitsProcessor(nn.Module):
             last_hidden = hidden_states[last_index]
 
         last_logits = torch.matmul(last_hidden, weight.T)
-        if self.tp_size > 1:
+        if self.do_tensor_parallel_all_gather:
             last_logits = tensor_model_parallel_all_gather(last_logits)
         last_logits = last_logits[:, : self.config.vocab_size].float()
 
         if hasattr(self.config, "final_logit_softcapping"):
-            last_logits /= self.config.final_logit_softcapping
-            last_logits = torch.tanh(last_logits)
-            last_logits *= self.config.final_logit_softcapping
+            last_logits.div_(self.config.final_logit_softcapping)
+            torch.tanh(last_logits, out=last_logits)
+            last_logits.mul_(self.config.final_logit_softcapping)
 
         # Return only last_logits if logprob is not requested
         if not logits_metadata.return_logprob:
-            return LogitProcessorOutput(
+            return LogitsProcessorOutput(
                 next_token_logits=last_logits,
                 next_token_logprobs=None,
                 normalized_prompt_logprobs=None,
@@ -194,7 +209,7 @@ class LogitsProcessor(nn.Module):
                 else:
                     output_top_logprobs = None
 
-                return LogitProcessorOutput(
+                return LogitsProcessorOutput(
                     next_token_logits=last_logits,
                     next_token_logprobs=last_logprobs,
                     normalized_prompt_logprobs=None,
@@ -203,15 +218,31 @@ class LogitsProcessor(nn.Module):
                     output_top_logprobs=output_top_logprobs,
                 )
             else:
-                all_logits = torch.matmul(hidden_states, weight.T)
-                if self.tp_size > 1:
+                pt, states, pruned_input_ids = 0, [], []
+                for i, extend_len in enumerate(logits_metadata.extend_seq_lens_cpu):
+                    start_len = logits_metadata.logprob_start_lens_cpu[i]
+                    states.append(hidden_states[pt + start_len : pt + extend_len])
+                    pruned_input_ids.append(input_ids[pt + start_len : pt + extend_len])
+                    pt += extend_len
+
+                states = torch.cat(states, dim=0)
+                pruned_input_ids = torch.cat(pruned_input_ids, dim=0)
+
+                cum_start_len1 = torch.tensor(
+                    logits_metadata.logprob_start_lens_cpu, device="cuda"
+                ).cumsum(0)
+                cum_start_len0 = torch.zeros_like(cum_start_len1)
+                cum_start_len0[1:] = cum_start_len1[:-1]
+
+                all_logits = torch.matmul(states, weight.T)
+                if self.do_tensor_parallel_all_gather:
                     all_logits = tensor_model_parallel_all_gather(all_logits)
                 all_logits = all_logits[:, : self.config.vocab_size].float()
 
                 if hasattr(self.config, "final_logit_softcapping"):
-                    all_logits /= self.config.final_logit_softcapping
-                    all_logits = torch.tanh(all_logits)
-                    all_logits *= self.config.final_logit_softcapping
+                    all_logits.div_(self.config.final_logit_softcapping)
+                    torch.tanh(all_logits, out=all_logits)
+                    all_logits.mul_(self.config.final_logit_softcapping)
 
                 all_logprobs = all_logits
                 del all_logits, hidden_states
@@ -228,20 +259,26 @@ class LogitsProcessor(nn.Module):
                 else:
                     input_top_logprobs = output_top_logprobs = None
 
-                last_logprobs = all_logprobs[last_index]
+                last_logprobs = all_logprobs[last_index - cum_start_len1]
 
                 # Compute the logprobs and normalized logprobs for the prefill tokens.
                 # Note that we pad a zero at the end of each sequence for easy computation.
                 input_token_logprobs = all_logprobs[
                     torch.arange(all_logprobs.shape[0], device="cuda"),
-                    torch.cat([input_ids[1:], torch.tensor([0], device="cuda")]),
+                    torch.cat([pruned_input_ids[1:], torch.tensor([0], device="cuda")]),
                 ]
 
                 normalized_prompt_logprobs = self._get_normalized_prompt_logprobs(
-                    input_token_logprobs, logits_metadata
+                    input_token_logprobs,
+                    cum_start_len0,
+                    cum_start_len1,
+                    logits_metadata,
                 )
 
-                return LogitProcessorOutput(
+                # Remove the last token logprob for the prefill tokens.
+                input_token_logprobs = input_token_logprobs[:-1]
+
+                return LogitsProcessorOutput(
                     next_token_logits=last_logits,
                     next_token_logprobs=last_logprobs,
                     normalized_prompt_logprobs=normalized_prompt_logprobs,

sglang/srt/layers/radix_attention.py

@@ -15,6 +15,8 @@ limitations under the License.
 
 """Radix attention."""
 
+from typing import Optional
+
 import torch
 from flashinfer.cascade import merge_state
 from torch import nn
@@ -34,6 +36,7 @@ class RadixAttention(nn.Module):
         scaling: float,
         num_kv_heads: int,
         layer_id: int,
+        sliding_window_size: Optional[int] = None,
         logit_cap: int = -1,
         v_head_dim: int = -1,
     ):
@@ -46,6 +49,7 @@ class RadixAttention(nn.Module):
         self.v_head_dim = v_head_dim if v_head_dim != -1 else head_dim
         self.scaling = scaling
         self.layer_id = layer_id
+        self.sliding_window_size = sliding_window_size if sliding_window_size else -1
 
         if (
             not global_server_args_dict.get("disable_flashinfer", False)
@@ -113,14 +117,25 @@ class RadixAttention(nn.Module):
         return o
 
     def extend_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
+        # using two wrappers is unnecessary in the current PR, but are prepared for future PRs
+        prefill_wrapper_paged = input_metadata.flashinfer_prefill_wrapper_paged
+        if self.sliding_window_size != -1:
+            prefill_wrapper_paged = prefill_wrapper_paged[0]
+        else:
+            if isinstance(prefill_wrapper_paged, list):
+                prefill_wrapper_paged = prefill_wrapper_paged[1]
+
         if not input_metadata.flashinfer_use_ragged:
-            self.store_kv_cache(k, v, input_metadata)
+            if k is not None:
+                assert v is not None
+                self.store_kv_cache(k, v, input_metadata)
 
-            o = input_metadata.flashinfer_prefill_wrapper_paged.forward(
+            o = prefill_wrapper_paged.forward(
                 q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
                 input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
                 causal=True,
                 sm_scale=self.scaling,
+                window_left=self.sliding_window_size,
                 logits_soft_cap=self.logit_cap,
             )
         else:
@@ -138,14 +153,12 @@ class RadixAttention(nn.Module):
             if input_metadata.extend_no_prefix:
                 o = o1
             else:
-                o2, s2 = (
-                    input_metadata.flashinfer_prefill_wrapper_paged.forward_return_lse(
-                        q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
-                        input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
-                        causal=False,
-                        sm_scale=self.scaling,
-                        logits_soft_cap=self.logit_cap,
-                    )
+                o2, s2 = prefill_wrapper_paged.forward_return_lse(
+                    q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
+                    input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
+                    causal=False,
+                    sm_scale=self.scaling,
+                    logits_soft_cap=self.logit_cap,
                 )
 
                 o, _ = merge_state(o1, s1, o2, s2)
@@ -158,9 +171,18 @@ class RadixAttention(nn.Module):
         return o.view(-1, self.tp_q_head_num * self.head_dim)
 
     def decode_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
-        self.store_kv_cache(k, v, input_metadata)
+        decode_wrapper = input_metadata.flashinfer_decode_wrapper
+        if self.sliding_window_size != -1:
+            decode_wrapper = decode_wrapper[0]
+        else:
+            if isinstance(decode_wrapper, list):
+                decode_wrapper = decode_wrapper[1]
 
-        o = input_metadata.flashinfer_decode_wrapper.forward(
+        if k is not None:
+            assert v is not None
+            self.store_kv_cache(k, v, input_metadata)
+
+        o = decode_wrapper.forward(
             q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
             input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
             sm_scale=self.scaling,
@@ -170,8 +192,10 @@ class RadixAttention(nn.Module):
         return o.view(-1, self.tp_q_head_num * self.head_dim)
 
     def forward(self, q, k, v, input_metadata: InputMetadata):
-        k = k.view(-1, self.tp_k_head_num, self.qk_head_dim)
-        v = v.view(-1, self.tp_v_head_num, self.v_head_dim)
+        if k is not None:
+            assert v is not None
+            k = k.view(-1, self.tp_k_head_num, self.qk_head_dim)
+            v = v.view(-1, self.tp_v_head_num, self.v_head_dim)
 
         if input_metadata.forward_mode == ForwardMode.EXTEND:
             return self.extend_forward(q, k, v, input_metadata)
@@ -179,7 +203,6 @@ class RadixAttention(nn.Module):
             return self.decode_forward(q, k, v, input_metadata)
 
     def store_kv_cache(self, cache_k, cache_v, input_metadata: InputMetadata):
-        k_cache = input_metadata.token_to_kv_pool.get_key_buffer(self.layer_id)
-        v_cache = input_metadata.token_to_kv_pool.get_value_buffer(self.layer_id)
-        k_cache[input_metadata.out_cache_loc] = cache_k
-        v_cache[input_metadata.out_cache_loc] = cache_v
+        input_metadata.token_to_kv_pool.set_kv_buffer(
+            self.layer_id, input_metadata.out_cache_loc, cache_k, cache_v
+        )

sglang/srt/layers/sampler.py

@@ -0,0 +1,154 @@
+import dataclasses
+import logging
+from typing import Union
+
+import torch
+from flashinfer.sampling import (
+    min_p_sampling_from_probs,
+    top_k_renorm_prob,
+    top_k_top_p_sampling_from_probs,
+    top_p_renorm_prob,
+)
+from vllm.model_executor.custom_op import CustomOp
+
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+
+# TODO: move this dict to another place
+from sglang.srt.managers.schedule_batch import global_server_args_dict
+from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class SampleOutput:
+    success: torch.Tensor
+    probs: torch.Tensor
+    batch_next_token_ids: torch.Tensor
+
+
+class Sampler(CustomOp):
+    def __init__(self):
+        super().__init__()
+
+    def _apply_penalties(self, logits: torch.Tensor, sampling_info: SamplingBatchInfo):
+        # min-token, presence, frequency
+        if sampling_info.linear_penalties is not None:
+            logits += sampling_info.linear_penalties
+
+        # repetition
+        if sampling_info.scaling_penalties is not None:
+            logits = torch.where(
+                logits > 0,
+                logits / sampling_info.scaling_penalties,
+                logits * sampling_info.scaling_penalties,
+            )
+
+        return logits
+
+    def _get_probs(
+        self,
+        logits: torch.Tensor,
+        sampling_info: SamplingBatchInfo,
+        is_torch_compile: bool = False,
+    ):
+        # Post process logits
+        logits = logits.contiguous()
+        logits.div_(sampling_info.temperatures)
+        if is_torch_compile:
+            # FIXME: Temporary workaround for unknown bugs in torch.compile
+            logits.add_(0)
+
+        if sampling_info.logit_bias is not None:
+            logits.add_(sampling_info.logit_bias)
+
+        if sampling_info.vocab_mask is not None:
+            logits = logits.masked_fill(~sampling_info.vocab_mask, float("-inf"))
+
+        logits = self._apply_penalties(logits, sampling_info)
+
+        return torch.softmax(logits, dim=-1)
+
+    def forward_cuda(
+        self,
+        logits: Union[torch.Tensor, LogitsProcessorOutput],
+        sampling_info: SamplingBatchInfo,
+    ):
+        if isinstance(logits, LogitsProcessorOutput):
+            logits = logits.next_token_logits
+
+        probs = self._get_probs(logits, sampling_info)
+
+        if not global_server_args_dict["disable_flashinfer_sampling"]:
+            max_top_k_round, batch_size = 32, probs.shape[0]
+            uniform_samples = torch.rand(
+                (max_top_k_round, batch_size), device=probs.device
+            )
+            if sampling_info.need_min_p_sampling:
+                probs = top_k_renorm_prob(probs, sampling_info.top_ks)
+                probs = top_p_renorm_prob(probs, sampling_info.top_ps)
+                batch_next_token_ids, success = min_p_sampling_from_probs(
+                    probs, uniform_samples, sampling_info.min_ps
+                )
+            else:
+                batch_next_token_ids, success = top_k_top_p_sampling_from_probs(
+                    probs, uniform_samples, sampling_info.top_ks, sampling_info.top_ps
+                )
+        else:
+            # Here we provide a slower fallback implementation.
+            batch_next_token_ids, success = top_k_top_p_min_p_sampling_from_probs_torch(
+                probs, sampling_info.top_ks, sampling_info.top_ps, sampling_info.min_ps
+            )
+
+        return SampleOutput(success, probs, batch_next_token_ids)
+
+    def forward_native(
+        self,
+        logits: Union[torch.Tensor, LogitsProcessorOutput],
+        sampling_info: SamplingBatchInfo,
+    ):
+        if isinstance(logits, LogitsProcessorOutput):
+            logits = logits.next_token_logits
+
+        probs = self._get_probs(logits, sampling_info, is_torch_compile=True)
+
+        batch_next_token_ids, success = top_k_top_p_min_p_sampling_from_probs_torch(
+            probs, sampling_info.top_ks, sampling_info.top_ps, sampling_info.min_ps
+        )
+
+        return SampleOutput(success, probs, batch_next_token_ids)
+
+
+def top_k_top_p_min_p_sampling_from_probs_torch(
+    probs: torch.Tensor,
+    top_ks: torch.Tensor,
+    top_ps: torch.Tensor,
+    min_ps: torch.Tensor,
+):
+    """A top-k, top-p and min-p sampling implementation with native pytorch operations."""
+    probs_sort, probs_idx = probs.sort(dim=-1, descending=True)
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    min_p_thresholds = probs_sort[:, 0] * min_ps
+    probs_sort[(probs_sum - probs_sort) > top_ps.view(-1, 1)] = 0.0
+    probs_sort[
+        torch.arange(0, probs.shape[-1], device=probs.device).view(1, -1)
+        >= top_ks.view(-1, 1)
+    ] = 0.0
+    probs_sort[probs_sort < min_p_thresholds.view(-1, 1)] = 0.0
+    probs_sort.div_(probs_sort.max(dim=-1, keepdim=True)[0])
+    try:
+        # FIXME: torch.multiomial does not support num_samples = 1
+        sampled_index = torch.multinomial(probs_sort, num_samples=2, replacement=True)[
+            :, :1
+        ]
+    except RuntimeError as e:
+        logger.warning(f"Sampling error: {e}")
+        batch_next_token_ids = torch.zeros(
+            (probs_sort.shape[0],), dtype=torch.int32, device=probs.device
+        )
+        success = torch.zeros(probs.shape[0], dtype=torch.bool, device=probs.device)
+        return batch_next_token_ids, success
+
+    batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index).view(-1)
+    success = torch.ones(probs.shape[0], dtype=torch.bool, device=probs.device)
+    return batch_next_token_ids, success

sglang/srt/managers/controller_multi.py

@@ -21,7 +21,6 @@ Each data parallel worker can manage multiple tensor parallel workers.
 import dataclasses
 import logging
 import multiprocessing
-import os
 from enum import Enum, auto
 
 import numpy as np
@@ -36,7 +35,7 @@ from sglang.srt.managers.io_struct import (
     TokenizedGenerateReqInput,
 )
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import kill_parent_process
+from sglang.srt.utils import configure_logger, kill_parent_process
 from sglang.utils import get_exception_traceback
 
 logger = logging.getLogger(__name__)
@@ -194,10 +193,7 @@ def start_controller_process(
 ):
     """Start a controller process."""
 
-    logging.basicConfig(
-        level=getattr(logging, server_args.log_level.upper()),
-        format="%(message)s",
-    )
+    configure_logger(server_args)
 
     try:
         controller = ControllerMulti(server_args, port_args, model_overide_args)
@@ -212,6 +208,4 @@ def start_controller_process(
     except Exception:
         logger.error("Exception in ControllerMulti:\n" + get_exception_traceback())
     finally:
-        for w in controller.workers:
-            os.kill(w.proc.pid, 9)
         kill_parent_process()

sglang/srt/managers/controller_single.py

@@ -17,7 +17,6 @@ limitations under the License.
 
 import logging
 import multiprocessing
-import os
 from typing import List
 
 import zmq
@@ -28,7 +27,7 @@ from sglang.srt.managers.tp_worker import (
     launch_tp_servers,
 )
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import kill_parent_process
+from sglang.srt.utils import configure_logger, kill_parent_process
 from sglang.utils import get_exception_traceback
 
 logger = logging.getLogger(__name__)
@@ -53,7 +52,7 @@ class ControllerSingle:
         self.dp_worker_id = dp_worker_id
         self.mp_queue = mp_queue
 
-        # Init communication
+        # Init inter-process communication
         context = zmq.Context(2)
 
         if not self.is_dp_worker:
@@ -134,11 +133,11 @@ def start_controller_process(
     queue: multiprocessing.connection.Connection = None,
 ):
     """Start a controller process."""
-
-    logging.basicConfig(
-        level=getattr(logging, server_args.log_level.upper()),
-        format="%(message)s",
-    )
+    if is_data_parallel_worker:
+        logger_prefix = f" DP{dp_worker_id} TP0"
+    else:
+        logger_prefix = " TP0"
+    configure_logger(server_args, prefix=logger_prefix)
 
     if not is_data_parallel_worker:
         tp_size_local = server_args.tp_size // server_args.nnodes
@@ -167,6 +166,4 @@ def start_controller_process(
     except Exception:
         logger.error("Exception in ControllerSingle:\n" + get_exception_traceback())
     finally:
-        for t in controller.tp_procs:
-            os.kill(t.pid, 9)
         kill_parent_process()

sglang/srt/managers/detokenizer_manager.py

@@ -17,7 +17,6 @@ limitations under the License.
 
 import asyncio
 import dataclasses
-import inspect
 from typing import List
 
 import uvloop
@@ -29,6 +28,7 @@ from sglang.srt.managers.io_struct import (
     BatchEmbeddingOut,
     BatchStrOut,
     BatchTokenIDOut,
+    UpdateWeightReqOutput,
 )
 from sglang.srt.managers.schedule_batch import FINISH_MATCHED_STR
 from sglang.srt.server_args import PortArgs, ServerArgs
@@ -39,6 +39,8 @@ asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 @dataclasses.dataclass
 class DecodeStatus:
+    """Store the status of incremental decoding."""
+
     vid: int
     decoded_text: str
     decode_ids: List[int]
@@ -47,11 +49,14 @@ class DecodeStatus:
 
 
 class DetokenizerManager:
+    """DetokenizerManager is a process that detokenizes the token ids."""
+
     def __init__(
         self,
         server_args: ServerArgs,
         port_args: PortArgs,
     ):
+        # Init inter-process communication
         context = zmq.asyncio.Context(2)
         self.recv_from_router = context.socket(zmq.PULL)
         self.recv_from_router.bind(f"tcp://127.0.0.1:{port_args.detokenizer_port}")
@@ -71,10 +76,13 @@ class DetokenizerManager:
         self.decode_status = {}
 
     async def handle_loop(self):
+        """The event loop that handles requests"""
+
         while True:
-            recv_obj: BatchTokenIDOut = await self.recv_from_router.recv_pyobj()
+            recv_obj = await self.recv_from_router.recv_pyobj()
 
             if isinstance(recv_obj, BatchEmbeddingOut):
+                # If it is embedding model, no detokenization is needed.
                 self.send_to_tokenizer.send_pyobj(
                     BatchEmbeddingOut(
                         rids=recv_obj.rids,
@@ -84,15 +92,18 @@ class DetokenizerManager:
                     )
                 )
                 continue
+            elif isinstance(recv_obj, UpdateWeightReqOutput):
+                # If it is a weight update request, no detokenization is needed.
+                self.send_to_tokenizer.send_pyobj(recv_obj)
+                continue
+            elif self.tokenizer is None:
+                # If the tokenizer is skipped, no detokenization is needed
+                self.send_to_tokenizer.send_pyobj(recv_obj)
+                continue
 
             assert isinstance(recv_obj, BatchTokenIDOut)
             bs = len(recv_obj.rids)
 
-            if self.tokenizer is None:
-                # Send BatchTokenIDOut if no tokenizer init'ed.
-                self.send_to_tokenizer.send_pyobj(recv_obj)
-                continue
-
             # Initialize decode status
             read_ids, surr_ids = [], []
             for i in range(bs):
@@ -126,8 +137,7 @@ class DetokenizerManager:
                 spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
             )
 
-            # Trim stop str
-            # TODO(lmzheng): handle the case where multiple stop strs are hit
+            # Incremental decoding
             output_strs = []
             for i in range(bs):
                 s = self.decode_status[recv_obj.rids[i]]
@@ -144,6 +154,7 @@ class DetokenizerManager:
 
                 output_strs.append(s.decoded_text + new_text)
 
+                # Trim stop str. TODO(lmzheng): handle the case where multiple stop strs are hit
                 if isinstance(recv_obj.finished_reason[i], FINISH_MATCHED_STR):
                     pos = output_strs[i].find(recv_obj.finished_reason[i].matched)
                     if pos != -1: