sglang 0.1.17__py3-none-any.whl → 0.1.18__py3-none-any.whl

sglang/srt/layers/logits_processor.py

@@ -1,3 +1,8 @@
+"""Logits processing."""
+
+import dataclasses
+from typing import List
+
 import torch
 from torch import nn
 from vllm.distributed import (
@@ -8,6 +13,24 @@ from vllm.distributed import (
 from sglang.srt.managers.controller.model_runner import ForwardMode, InputMetadata
 
 
+@dataclasses.dataclass
+class LogitProcessorOutput:
+    # 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]
+    next_token_logprobs: torch.Tensor
+
+    # The normlaized logprobs of prompts.  shape: [#seq]
+    normalized_prompt_logprobs: torch.Tensor
+    # The logprobs of prefill tokens.      shape: [#token, vocab_size]
+    prefill_token_logprobs: torch.Tensor
+
+    # The logprob and id of the top-k tokens in prefill positions.  shape [#seq, #token, k] of Tuple(logprob, token_id)
+    prefill_top_logprobs: List
+    # The logprob and id of the top-k tokens in decode positions.   shape [#seq, #token, k] of Tuple(logprob, token_id)
+    decode_top_logprobs: List
+
+
 class LogitsProcessor(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -37,6 +60,7 @@ class LogitsProcessor(nn.Module):
         return normalized_prompt_logprobs
 
     def _get_top_logprobs(self, all_logprobs, input_metadata: InputMetadata):
+        # TODO: vectorize the code below
         if input_metadata.forward_mode == ForwardMode.DECODE:
             decode_top_logprobs = []
             for i in range(all_logprobs.shape[0]):
@@ -49,7 +73,6 @@ class LogitsProcessor(nn.Module):
         else:
             prefill_top_logprobs, decode_top_logprobs = [], []
             pt = 0
-            # NOTE: the GPU-CPU overhead can be reduced
             extend_seq_lens_cpu = input_metadata.extend_seq_lens.tolist()
             for i, extend_seq_len in enumerate(extend_seq_lens_cpu):
                 if extend_seq_len == 0:
@@ -69,18 +92,15 @@ class LogitsProcessor(nn.Module):
             return prefill_top_logprobs, decode_top_logprobs
 
     def forward(self, input_ids, hidden_states, weight, input_metadata: InputMetadata):
-        # Get last index for next token prediction, except for DECODE mode.
-        last_index = None
-        if input_metadata.forward_mode != ForwardMode.DECODE:
+        # Get the last hidden states and last logits for the next token prediction
+        if input_metadata.forward_mode == ForwardMode.DECODE:
+            last_index = None
+            last_hidden = hidden_states
+        else:
             last_index = (
                 torch.cumsum(input_metadata.extend_seq_lens, dim=0, dtype=torch.long)
                 - 1
             )
-
-        # Get the last hidden states and last logits
-        if input_metadata.forward_mode == ForwardMode.DECODE:
-            last_hidden = hidden_states
-        else:
             last_hidden = hidden_states[last_index]
 
         last_logits = torch.matmul(last_hidden, weight.T)
@@ -90,8 +110,14 @@ class LogitsProcessor(nn.Module):
 
         # Return only last_logits if logprob is not requested
         if not input_metadata.return_logprob:
-            hidden_states = None
-            return last_logits, (None, None, None, None, None)
+            return LogitProcessorOutput(
+                next_token_logits=last_logits,
+                next_token_logprobs=None,
+                normalized_prompt_logprobs=None,
+                prefill_token_logprobs=None,
+                prefill_top_logprobs=None,
+                decode_top_logprobs=None,
+            )
         else:
             # When logprob is requested, compute the logits for all tokens.
             if input_metadata.forward_mode == ForwardMode.DECODE:
@@ -106,6 +132,7 @@ class LogitsProcessor(nn.Module):
             del all_logits
             all_logprobs[:] = torch.nn.functional.log_softmax(all_logprobs, dim=-1)
 
+            # Get the logprob of top-k tokens
             return_top_logprob = any(x > 0 for x in input_metadata.top_logprobs_nums)
             if return_top_logprob:
                 prefill_top_logprobs, decode_top_logprobs = self._get_top_logprobs(
@@ -115,16 +142,15 @@ class LogitsProcessor(nn.Module):
                 prefill_top_logprobs = decode_top_logprobs = None
 
             if input_metadata.forward_mode == ForwardMode.DECODE:
-                last_logprobs = all_logprobs
-                return last_logits, (
-                    None,
-                    None,
-                    None,
-                    decode_top_logprobs,
-                    last_logprobs,
+                return LogitProcessorOutput(
+                    next_token_logits=last_logits,
+                    next_token_logprobs=all_logprobs,
+                    normalized_prompt_logprobs=None,
+                    prefill_token_logprobs=None,
+                    prefill_top_logprobs=None,
+                    decode_top_logprobs=decode_top_logprobs,
                 )
             else:
-                # Compute the logprobs for the last token of each request.
                 last_logprobs = all_logprobs[last_index]
 
                 # Compute the logprobs and normalized logprobs for the prefill tokens.
@@ -137,12 +163,14 @@ class LogitsProcessor(nn.Module):
                 normalized_prompt_logprobs = self._get_normalized_prompt_logprobs(
                     prefill_token_logprobs, input_metadata
                 )
-                return last_logits, (
-                    prefill_token_logprobs,
-                    normalized_prompt_logprobs,
-                    prefill_top_logprobs,
-                    decode_top_logprobs,
-                    last_logprobs,
+
+                return LogitProcessorOutput(
+                    next_token_logits=last_logits,
+                    next_token_logprobs=last_logprobs,
+                    normalized_prompt_logprobs=normalized_prompt_logprobs,
+                    prefill_token_logprobs=prefill_token_logprobs,
+                    prefill_top_logprobs=prefill_top_logprobs,
+                    decode_top_logprobs=decode_top_logprobs,
                 )
 
 

sglang/srt/layers/radix_attention.py

@@ -1,7 +1,10 @@
-import torch
+"""Radix attention."""
+
 import numpy as np
+import torch
 from torch import nn
 
+from sglang.global_config import global_config
 from sglang.srt.layers.context_flashattention_nopad import context_attention_fwd
 from sglang.srt.layers.extend_attention import extend_attention_fwd
 from sglang.srt.layers.token_attention import token_attention_fwd
@@ -9,27 +12,32 @@ from sglang.srt.managers.controller.model_runner import ForwardMode, InputMetada
 
 
 class RadixAttention(nn.Module):
-    def __init__(self, num_heads, head_dim, scaling, num_kv_heads, layer_id, logit_cap=-1):
+    def __init__(
+        self, num_heads: int, head_dim: int, scaling: float, num_kv_heads: int,
+        layer_id: int, logit_cap: int = -1
+    ):
         super().__init__()
         self.tp_q_head_num = num_heads
         self.tp_k_head_num = num_kv_heads
         self.tp_v_head_num = num_kv_heads
         self.head_dim = head_dim
         self.layer_id = layer_id
-        self.logit_cap = logit_cap
 
         assert np.allclose(scaling, 1.0 / (head_dim**0.5))
 
         from sglang.srt.managers.controller.model_runner import global_server_args_dict
 
-        if global_server_args_dict.get("enable_flashinfer", False):
+        if not global_server_args_dict.get("disable_flashinfer", False):
             self.prefill_forward = self.prefill_forward_flashinfer
             self.extend_forward = self.prefill_forward_flashinfer
             self.decode_forward = self.decode_forward_flashinfer
+            # flashinfer now accepts float logit_cap argument
+            self.logit_cap = logit_cap if logit_cap > 0 else 0
         else:
             self.prefill_forward = self.prefill_forward_triton
             self.extend_forward = self.extend_forward_triton
             self.decode_forward = self.decode_forward_triton
+            self.logit_cap = logit_cap
 
     def prefill_forward_triton(self, q, k, v, input_metadata: InputMetadata):
         o = torch.empty_like(q)
@@ -96,19 +104,38 @@ class RadixAttention(nn.Module):
     def prefill_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
         self.store_kv_cache(k, v, input_metadata)
 
-        o = input_metadata.prefill_wrapper.forward(
+        o1, s1 = input_metadata.flashinfer_prefill_wrapper_ragged.forward_return_lse(
             q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
-            input_metadata.token_to_kv_pool.kv_data[self.layer_id],
+            k.contiguous().view(-1, self.tp_k_head_num, self.head_dim),
+            v.contiguous().view(-1, self.tp_v_head_num, self.head_dim),
+            logits_soft_cap=self.logit_cap,
         )
 
+        if input_metadata.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.kv_data[self.layer_id],
+                causal=False,
+                logits_soft_cap=self.logit_cap,
+            )
+
+            from flashinfer.cascade import merge_state
+            o, _ = merge_state(o1, s1, o2, s2)
+
+        if input_metadata.total_num_tokens >= global_config.layer_sync_threshold:
+            torch.cuda.synchronize()
+
         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)
 
-        o = input_metadata.decode_wrapper.forward(
+        o = input_metadata.flashinfer_decode_wrapper.forward(
             q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
             input_metadata.token_to_kv_pool.kv_data[self.layer_id],
+            logits_soft_cap=self.logit_cap,
         )
 
         return o.view(-1, self.tp_q_head_num * self.head_dim)

sglang/srt/managers/controller/dp_worker.py

@@ -1,9 +1,10 @@
 """A data parallel worker thread."""
+
 import asyncio
 import logging
 import queue
 import threading
-from typing import List, Callable
+from typing import Callable, List
 
 import uvloop
 import zmq
@@ -69,7 +70,9 @@ class DataParallelWorkerThread(threading.Thread):
 
             # async sleep for receiving the subsequent request and avoiding cache miss
             if len(out_pyobjs) != 0:
-                has_finished = any([obj.finished_reason is not None for obj in out_pyobjs])
+                has_finished = any(
+                    [obj.finished_reason is not None for obj in out_pyobjs]
+                )
                 if has_finished:
                     await asyncio.sleep(self.request_dependency_delay)
             await asyncio.sleep(global_config.wait_for_new_request_delay)
@@ -107,4 +110,4 @@ def start_data_parallel_worker(
         step_func=model_tp_client.step,
     )
     worker_thread.start()
-    return worker_thread
+    return worker_thread

sglang/srt/managers/controller/infer_batch.py

@@ -1,4 +1,6 @@
 """Meta data for requests and batches"""
+
+import warnings
 from dataclasses import dataclass
 from enum import IntEnum, auto
 from typing import List
@@ -6,9 +8,13 @@ from typing import List
 import numpy as np
 import torch
 
+from sglang.srt.constrained import RegexGuide
+from sglang.srt.constrained.jump_forward import JumpForwardMap
 from sglang.srt.managers.controller.radix_cache import RadixCache
 from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
 
+INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5
+
 
 class ForwardMode(IntEnum):
     PREFILL = auto()
@@ -63,12 +69,15 @@ class Req:
     def __init__(self, rid, origin_input_text, origin_input_ids):
         self.rid = rid
         self.origin_input_text = origin_input_text
+        self.origin_input_ids_unpadded = origin_input_ids  # Before image padding
         self.origin_input_ids = origin_input_ids
-        self.origin_input_ids_unpadded = origin_input_ids  # before image padding
-        self.prev_output_str = ""
-        self.prev_output_ids = []
-        self.output_ids = []
-        self.input_ids = None  # input_ids = origin_input_ids + prev_output_ids
+        self.output_ids = []  # Each decode stage's output ids
+        self.input_ids = None  # input_ids = origin_input_ids + output_ids
+
+        # For incremental decode
+        self.decoded_text = ""
+        self.surr_offset = None  # Surrounding offset to defeat the cleanup algorithm
+        self.read_offset = None
 
         # The number of decoded tokens for token usage report. Note that
         # this does not include the jump forward tokens.
@@ -108,20 +117,54 @@ class Req:
         self.last_update_decode_tokens = 0
 
         # Constrained decoding
-        self.regex_fsm = None
-        self.regex_fsm_state = 0
-        self.jump_forward_map = None
+        self.regex_fsm: RegexGuide = None
+        self.regex_fsm_state: int = 0
+        self.jump_forward_map: JumpForwardMap = None
 
     # whether request reached finished condition
     def finished(self) -> bool:
         return self.finished_reason is not None
 
-    def partial_decode(self, ids):
-        first_token = self.tokenizer.convert_ids_to_tokens(ids[0])
-        first_token = (
-            first_token.decode() if isinstance(first_token, bytes) else first_token
+    # Based on https://github.com/vllm-project/vllm/blob/7a64d24aad69e4d2548aa0bf528d9fe63428ab01/vllm/transformers_utils/detokenizer.py#L194-L313
+    def init_detokenize_incrementally(self):
+        first_iter = self.surr_offset is None or self.read_offset is None
+
+        if first_iter:
+            self.read_offset = len(self.origin_input_ids_unpadded)
+            self.surr_offset = max(
+                self.read_offset - INIT_INCREMENTAL_DETOKENIZATION_OFFSET, 0
+            )
+
+        all_ids = self.origin_input_ids_unpadded + self.output_ids
+        surr_ids = all_ids[self.surr_offset : self.read_offset]
+        read_ids = all_ids[self.surr_offset :]
+
+        return surr_ids, read_ids, len(all_ids)
+
+    def detokenize_incrementally(self, inplace: bool = True):
+        surr_ids, read_ids, num_all_tokens = self.init_detokenize_incrementally()
+
+        surr_text = self.tokenizer.decode(
+            surr_ids,
+            skip_special_tokens=self.sampling_params.skip_special_tokens,
+            spaces_between_special_tokens=self.sampling_params.spaces_between_special_tokens,
+        )
+        new_text = self.tokenizer.decode(
+            read_ids,
+            skip_special_tokens=self.sampling_params.skip_special_tokens,
+            spaces_between_special_tokens=self.sampling_params.spaces_between_special_tokens,
         )
-        return (" " if first_token.startswith("▁") else "") + self.tokenizer.decode(ids)
+
+        if len(new_text) > len(surr_text) and not new_text.endswith("�"):
+            new_text = new_text[len(surr_text) :]
+            if inplace:
+                self.decoded_text += new_text
+                self.surr_offset = self.read_offset
+                self.read_offset = num_all_tokens
+
+            return True, new_text
+
+        return False, ""
 
     def max_new_tokens(self):
         return self.sampling_params.max_new_tokens
@@ -130,18 +173,17 @@ class Req:
         if self.finished():
             return
 
-        if (
-            len(self.prev_output_ids) + len(self.output_ids)
-            >= self.sampling_params.max_new_tokens
-        ):
-            self.finished_reason = FINISH_LENGTH(len(self.prev_output_ids) + len(self.output_ids))
+        if len(self.output_ids) >= self.sampling_params.max_new_tokens:
+            self.finished_reason = FINISH_LENGTH(len(self.output_ids))
             return
 
         if (
             self.output_ids[-1] == self.tokenizer.eos_token_id
             and not self.sampling_params.ignore_eos
         ):
-            self.finished_reason = FINISH_MATCHED_TOKEN(matched=self.tokenizer.eos_token_id)
+            self.finished_reason = FINISH_MATCHED_TOKEN(
+                matched=self.tokenizer.eos_token_id
+            )
             return
 
         if len(self.sampling_params.stop_strs) > 0:
@@ -150,61 +192,59 @@ class Req:
             )
 
             for stop_str in self.sampling_params.stop_strs:
-                # FIXME: (minor) try incremental match in prev_output_str
-                if stop_str in tail_str or stop_str in self.prev_output_str:
+                if stop_str in tail_str or stop_str in self.decoded_text:
                     self.finished_reason = FINISH_MATCHED_STR(matched=stop_str)
                     return
 
     def jump_forward_and_retokenize(self, jump_forward_str, next_state):
-        # FIXME: This logic does not really solve the problem of determining whether
-        # there should be a leading space.
-        cur_output_str = self.partial_decode(self.output_ids)
-
-        # TODO(lsyin): apply re-tokenize only for decode tokens so that we do not need origin_input_text anymore
         if self.origin_input_text is None:
             # Recovering text can only use unpadded ids
             self.origin_input_text = self.tokenizer.decode(
                 self.origin_input_ids_unpadded
             )
 
-        all_text = (
-            self.origin_input_text
-            + self.prev_output_str
-            + cur_output_str
-            + jump_forward_str
-        )
+        all_text = self.origin_input_text + self.decoded_text + jump_forward_str
         all_ids = self.tokenizer.encode(all_text)
         prompt_tokens = len(self.origin_input_ids_unpadded)
-        self.origin_input_ids = all_ids[:prompt_tokens]
-        self.origin_input_ids_unpadded = self.origin_input_ids
-        # NOTE: the output ids may not strictly correspond to the output text
-        old_prev_output_ids = self.prev_output_ids
-        self.prev_output_ids = all_ids[prompt_tokens:]
-        self.prev_output_str = self.prev_output_str + cur_output_str + jump_forward_str
-        self.output_ids = []
+
+        if all_ids[prompt_tokens - 1] != self.origin_input_ids_unpadded[-1]:
+            # TODO(lsyin): fix token fusion
+            warnings.warn(
+                "Token fusion between input and output, try to avoid this by removing the space at the end of the input."
+            )
+            return False
+
+        old_output_ids = self.output_ids
+        self.output_ids = all_ids[prompt_tokens:]
+        self.decoded_text = self.decoded_text + jump_forward_str
+        self.surr_offset = prompt_tokens
+        self.read_offset = len(all_ids)
+
+        # NOTE: A trick to reduce the surrouding tokens decoding overhead
+        for i in range(0, INIT_INCREMENTAL_DETOKENIZATION_OFFSET):
+            surr_text_ = self.tokenizer.decode(
+                all_ids[self.read_offset - i : self.read_offset]
+            )
+            if not surr_text_.endswith("�"):
+                self.surr_offset = self.read_offset - i
+                break
 
         self.regex_fsm_state = next_state
 
         if self.return_logprob:
             # For fast-forward part's logprobs
             k = 0
-            for i, old_id in enumerate(old_prev_output_ids):
-                if old_id == self.prev_output_ids[i]:
+            for i, old_id in enumerate(old_output_ids):
+                if old_id == self.output_ids[i]:
                     k = k + 1
                 else:
                     break
             self.decode_token_logprobs = self.decode_token_logprobs[:k]
             self.decode_top_logprobs = self.decode_top_logprobs[:k]
             self.logprob_start_len = prompt_tokens + k
-            self.last_update_decode_tokens = len(self.prev_output_ids) - k
-
-        # print("=" * 100)
-        # print(f"Catch jump forward:\n{jump_forward_str}")
-        # print(self.tokenizer.convert_ids_to_tokens(self.input_ids))
-        # print(self.tokenizer.convert_ids_to_tokens(new_input_ids))
+            self.last_update_decode_tokens = len(self.output_ids) - k
 
-        # print(f"Output and jump forward str:\n{self.output_and_jump_forward_str}")
-        # print("*" * 100)
+        return True
 
     def __repr__(self):
         return f"rid(n={self.rid}, " f"input_ids={self.origin_input_ids}, "
@@ -263,6 +303,10 @@ class Batch:
     def is_empty(self):
         return len(self.reqs) == 0
 
+    # whether batch has at least 1 streaming request
+    def has_stream(self) -> bool:
+        return any(r.stream for r in self.reqs)
+
     def prepare_for_extend(self, vocab_size: int, int_token_logit_bias: torch.Tensor):
         device = "cuda"
         bs = len(self.reqs)
@@ -380,7 +424,10 @@ class Batch:
         sorted_indices = [i for i in range(len(self.reqs))]
         # TODO(lsyin): improve the priority of retraction
         sorted_indices.sort(
-            key=lambda i: (len(self.reqs[i].output_ids), -len(self.reqs[i].input_ids)),
+            key=lambda i: (
+                len(self.reqs[i].output_ids),
+                -len(self.reqs[i].origin_input_ids),
+            ),
             reverse=True,
         )
 
@@ -402,14 +449,9 @@ class Batch:
             # release the last node
             self.tree_cache.dec_lock_ref(req.last_node)
 
-            cur_output_str = req.partial_decode(req.output_ids)
-            req.prev_output_str = req.prev_output_str + cur_output_str
-            req.prev_output_ids.extend(req.output_ids)
-
             req.prefix_indices = None
             req.last_node = None
             req.extend_input_len = 0
-            req.output_ids = []
 
             # For incremental logprobs
             req.last_update_decode_tokens = 0
@@ -427,18 +469,54 @@ class Batch:
 
         for i, req in enumerate(self.reqs):
             if req.jump_forward_map is not None:
-                res = req.jump_forward_map.jump_forward(req.regex_fsm_state)
-                if res is not None:
-                    jump_forward_str, next_state = res
-                    if len(jump_forward_str) <= 1:
+                jump_forward_bytes = req.jump_forward_map.jump_forward_byte(
+                    req.regex_fsm_state
+                )
+                if jump_forward_bytes is not None and len(jump_forward_bytes) > 1:
+                    suffix_bytes = []
+                    continuation_range = range(0x80, 0xC0)
+                    cur_state = req.regex_fsm_state
+                    while (
+                        len(jump_forward_bytes)
+                        and jump_forward_bytes[0][0] in continuation_range
+                    ):
+                        # continuation bytes
+                        byte_edge = jump_forward_bytes.pop(0)
+                        suffix_bytes.append(byte_edge[0])
+                        cur_state = byte_edge[1]
+
+                    suffix_tokens = [f"<0x{hex(b)[2:].upper()}>" for b in suffix_bytes]
+                    suffix_ids = req.tokenizer.convert_tokens_to_ids(suffix_tokens)
+
+                    # Current ids, for cache and revert
+                    cur_all_ids = tuple(req.origin_input_ids + req.output_ids)[:-1]
+                    cur_output_ids = req.output_ids
+
+                    req.output_ids.extend(suffix_ids)
+                    decode_res, new_text = req.detokenize_incrementally(inplace=False)
+                    if not decode_res:
+                        req.output_ids = cur_output_ids
                         continue
 
-                    if req_pool_indices_cpu is None:
-                        req_pool_indices_cpu = self.req_pool_indices.tolist()
+                    (
+                        jump_forward_str,
+                        next_state,
+                    ) = req.jump_forward_map.jump_forward_symbol(cur_state)
+
+                    # Make the incrementally decoded text part of jump_forward_str
+                    # so that the UTF-8 will not corrupt
+                    jump_forward_str = new_text + jump_forward_str
+                    if not req.jump_forward_and_retokenize(
+                        jump_forward_str, next_state
+                    ):
+                        req.output_ids = cur_output_ids
+                        continue
 
                     # insert the old request into tree_cache
+                    if req_pool_indices_cpu is None:
+                        req_pool_indices_cpu = self.req_pool_indices.tolist()
                     self.tree_cache.cache_req(
-                        token_ids=tuple(req.input_ids + req.output_ids)[:-1],
+                        token_ids=cur_all_ids,
                         last_uncached_pos=len(req.prefix_indices),
                         req_pool_idx=req_pool_indices_cpu[i],
                     )
@@ -446,9 +524,6 @@ class Batch:
                     # unlock the last node
                     self.tree_cache.dec_lock_ref(req.last_node)
 
-                    # jump-forward
-                    req.jump_forward_and_retokenize(jump_forward_str, next_state)
-
                     # re-applying image padding
                     if req.pixel_values is not None:
                         (
@@ -582,7 +657,7 @@ class Batch:
                 if req.regex_fsm is not None:
                     allowed_mask.zero_()
                     allowed_mask[
-                        req.regex_fsm.allowed_token_ids(req.regex_fsm_state)
+                        req.regex_fsm.get_next_instruction(req.regex_fsm_state).tokens
                     ] = 1
                     logits[i].masked_fill_(~allowed_mask, float("-inf"))
 
@@ -601,7 +676,7 @@ class Batch:
             batch_next_token_ids_cpu = batch_next_token_ids.cpu().numpy()
             for i, req in enumerate(self.reqs):
                 if req.regex_fsm is not None:
-                    req.regex_fsm_state = req.regex_fsm.next_state(
+                    req.regex_fsm_state = req.regex_fsm.get_next_state(
                         req.regex_fsm_state, batch_next_token_ids_cpu[i]
                     )
 

sglang/srt/managers/controller/manager_multi.py

@@ -13,15 +13,15 @@ import zmq
 import zmq.asyncio
 
 from sglang.global_config import global_config
+from sglang.srt.managers.controller.dp_worker import (
+    DataParallelWorkerThread,
+    start_data_parallel_worker,
+)
 from sglang.srt.managers.io_struct import (
     AbortReq,
     FlushCacheReq,
     TokenizedGenerateReqInput,
 )
-from sglang.srt.managers.controller.dp_worker import (
-    DataParallelWorkerThread,
-    start_data_parallel_worker,
-)
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.utils import get_exception_traceback
 
@@ -136,7 +136,7 @@ class Controller:
                 self.recv_reqs = []
                 if next_step_input:
                     await self.dispatching(next_step_input)
-            #else:
+            # else:
             #    logger.error("There is no live worker.")
 
             await asyncio.sleep(global_config.wait_for_new_request_delay)

sglang/srt/managers/controller/manager_single.py

@@ -1,7 +1,8 @@
 """A controller that manages a group of tensor parallel workers."""
+
 import asyncio
 import logging
-import time
+from concurrent.futures import ThreadPoolExecutor
 
 import uvloop
 import zmq
@@ -49,7 +50,9 @@ class ControllerSingle:
             # async sleep for receiving the subsequent request and avoiding cache miss
             slept = False
             if len(out_pyobjs) != 0:
-                has_finished = any([obj.finished_reason is not None for obj in out_pyobjs])
+                has_finished = any(
+                    [obj.finished_reason is not None for obj in out_pyobjs]
+                )
                 if has_finished:
                     if self.request_dependency_delay > 0:
                         slept = True
@@ -73,8 +76,9 @@ def start_controller_process(
     )
 
     try:
+        tp_size_local = server_args.tp_size // server_args.nnodes
         model_client = ModelTpClient(
-            list(range(server_args.tp_size)),
+            [i for _ in range(server_args.nnodes) for i in range(tp_size_local)],
             server_args,
             port_args.model_port_args[0],
             model_overide_args,
@@ -87,6 +91,7 @@ def start_controller_process(
     pipe_writer.send("init ok")
 
     loop = asyncio.new_event_loop()
+    loop.set_default_executor(ThreadPoolExecutor(max_workers=256))
     asyncio.set_event_loop(loop)
     loop.create_task(controller.loop_for_recv_requests())
     try: