sglang 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

sglang/srt/managers/{controller/infer_batch.py → schedule_batch.py}

@@ -1,5 +1,21 @@
+"""
+Copyright 2023-2024 SGLang Team
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
 """Meta data for requests and batches"""
 
+import logging
 import warnings
 from dataclasses import dataclass
 from enum import IntEnum, auto
@@ -12,11 +28,21 @@ from flashinfer.sampling import top_k_top_p_sampling_from_probs
 from sglang.global_config import global_config
 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
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool, TokenToKVPool
+from sglang.srt.mem_cache.radix_cache import RadixCache
 
 INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5
 
+# Put some global args for easy access
+global_server_args_dict = {
+    "disable_flashinfer": False,
+    "disable_flashinfer_sampling": False,
+    "attention_reduce_in_fp32": False,
+}
+
+
+logger = logging.getLogger(__name__)
+
 
 class ForwardMode(IntEnum):
     # Prefill a new sequence. This is deprecated now. "EXTEND" covers this case.
@@ -124,10 +150,10 @@ class Req:
         self.logprob_start_len = 0
         self.top_logprobs_num = 0
         self.normalized_prompt_logprob = None
-        self.prefill_token_logprobs = None
-        self.prefill_top_logprobs = None
-        self.decode_token_logprobs = []
-        self.decode_top_logprobs = []
+        self.input_token_logprobs = None
+        self.input_top_logprobs = None
+        self.output_token_logprobs = []
+        self.output_top_logprobs = []
         # The tokens is prefilled but need to be considered as decode tokens
         # and should be updated for the decode logprobs
         self.last_update_decode_tokens = 0
@@ -244,8 +270,8 @@ class Req:
                     k = k + 1
                 else:
                     break
-            self.decode_token_logprobs = self.decode_token_logprobs[:k]
-            self.decode_top_logprobs = self.decode_top_logprobs[:k]
+            self.output_token_logprobs = self.output_token_logprobs[:k]
+            self.output_top_logprobs = self.output_top_logprobs[:k]
             self.logprob_start_len = prompt_tokens + k
             self.last_update_decode_tokens = len(self.output_ids) - k
 
@@ -357,7 +383,7 @@ class Batch:
             out_cache_loc = self.token_to_kv_pool.alloc(extend_num_tokens)
 
             if out_cache_loc is None:
-                print("Prefill out of memory. This should never happen.")
+                logger.error("Prefill out of memory. This should never happen.")
                 self.tree_cache.pretty_print()
                 exit()
 
@@ -376,7 +402,7 @@ class Batch:
                     logit_bias = torch.zeros(
                         (bs, vocab_size), dtype=torch.float32, device=device
                     )
-                logit_bias[i] = int_token_logit_bias
+                logit_bias[i][: len(int_token_logit_bias)] = int_token_logit_bias
 
         # Set fields
         self.input_ids = torch.tensor(
@@ -591,7 +617,7 @@ class Batch:
         self.out_cache_loc = self.token_to_kv_pool.alloc(bs)
 
         if self.out_cache_loc is None:
-            print("Decode out of memory. This should never happen.")
+            logger.error("Decode out of memory. This should never happen.")
             self.tree_cache.pretty_print()
             exit()
 
@@ -687,13 +713,21 @@ class Batch:
         # TODO(lmzheng): apply penalty
         probs = torch.softmax(logits, dim=-1)
 
-        max_top_k_round, batch_size = 32, probs.shape[0]
-        uniform_samples = torch.rand((max_top_k_round, batch_size), device=probs.device)
-        batch_next_token_ids, success = top_k_top_p_sampling_from_probs(
-            probs, uniform_samples, self.top_ks, self.top_ps
-        )
+        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
+            )
+            batch_next_token_ids, success = top_k_top_p_sampling_from_probs(
+                probs, uniform_samples, self.top_ks, self.top_ps
+            )
+        else:
+            # Here we provide a slower fallback implementation.
+            batch_next_token_ids, success = top_k_top_p_sampling_from_probs_torch(
+                probs, self.top_ks, self.top_ps
+            )
 
-        if torch.any(~success):
+        if not torch.all(success):
             warnings.warn("Sampling failed, fallback to top_k=1 strategy")
             probs = probs.masked_fill(torch.isnan(probs), 0.0)
             argmax_ids = torch.argmax(probs, dim=-1)
@@ -747,7 +781,7 @@ class InputMetadata:
     flashinfer_prefill_wrapper_ragged: "BatchPrefillWithRaggedKVCacheWrapper" = None
     flashinfer_prefill_wrapper_paged: "BatchPrefillWithPagedKVCacheWrapper" = None
     flashinfer_decode_wrapper: "BatchDecodeWithPagedKVCacheWrapper" = None
-    use_ragged: bool = False
+    flashinfer_use_ragged: bool = False
 
     @classmethod
     def create(
@@ -763,10 +797,10 @@ class InputMetadata:
         return_logprob=False,
         skip_flashinfer_init=False,
     ):
-        use_ragged = False
+        flashinfer_use_ragged = False
         if not skip_flashinfer_init and not model_runner.server_args.disable_flashinfer:
             if forward_mode != ForwardMode.DECODE and int(torch.sum(seq_lens)) > 4096:
-                use_ragged = True
+                flashinfer_use_ragged = True
             init_flashinfer_args(
                 forward_mode,
                 model_runner,
@@ -774,7 +808,7 @@ class InputMetadata:
                 seq_lens,
                 prefix_lens,
                 model_runner.flashinfer_decode_wrapper,
-                use_ragged,
+                flashinfer_use_ragged,
             )
 
         batch_size = len(req_pool_indices)
@@ -829,7 +863,7 @@ class InputMetadata:
             flashinfer_prefill_wrapper_ragged=model_runner.flashinfer_prefill_wrapper_ragged,
             flashinfer_prefill_wrapper_paged=model_runner.flashinfer_prefill_wrapper_paged,
             flashinfer_decode_wrapper=model_runner.flashinfer_decode_wrapper,
-            use_ragged=use_ragged,
+            flashinfer_use_ragged=flashinfer_use_ragged,
         )
 
         if model_runner.server_args.disable_flashinfer:
@@ -850,7 +884,7 @@ def init_flashinfer_args(
     seq_lens,
     prefix_lens,
     flashinfer_decode_wrapper,
-    use_ragged=False,
+    flashinfer_use_ragged=False,
 ):
     """Init auxiliary variables for FlashInfer attention backend."""
     num_qo_heads = model_runner.model_config.num_attention_heads // model_runner.tp_size
@@ -859,7 +893,7 @@ def init_flashinfer_args(
     batch_size = len(req_pool_indices)
     total_num_tokens = int(torch.sum(seq_lens))
 
-    if use_ragged:
+    if flashinfer_use_ragged:
         paged_kernel_lens = prefix_lens
     else:
         paged_kernel_lens = seq_lens
@@ -895,7 +929,7 @@ def init_flashinfer_args(
         qo_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda")
         qo_indptr[1:] = torch.cumsum(seq_lens - prefix_lens, dim=0)
 
-        if use_ragged:
+        if flashinfer_use_ragged:
             model_runner.flashinfer_prefill_wrapper_ragged.end_forward()
             model_runner.flashinfer_prefill_wrapper_ragged.begin_forward(
                 qo_indptr,
@@ -933,3 +967,29 @@ def init_triton_args(forward_mode, seq_lens, prefix_lens):
         max_extend_len = int(torch.max(extend_seq_lens))
 
     return max_seq_len, max_extend_len, start_loc, prefix_lens
+
+
+def top_k_top_p_sampling_from_probs_torch(
+    probs: torch.Tensor, top_ks: torch.Tensor, top_ps: torch.Tensor
+):
+    """A top-k and top-k sampling implementation with native pytorch operations."""
+    probs_sort, probs_idx = probs.sort(dim=-1, descending=True)
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    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.div_(probs_sort.max(dim=-1, keepdim=True)[0])
+    try:
+        sampled_index = torch.multinomial(probs_sort, num_samples=1)
+    except RuntimeError:
+        batch_next_token_ids = torch.zeros(
+            (probs_sort.shape[0],), dtype=torch.int64, 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/tokenizer_manager.py

@@ -1,3 +1,18 @@
+"""
+Copyright 2023-2024 SGLang Team
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
 """TokenizerManager is a process that tokenizes the text."""
 
 import asyncio
@@ -6,7 +21,7 @@ import dataclasses
 import logging
 import multiprocessing as mp
 import os
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 import numpy as np
 import transformers
@@ -69,6 +84,7 @@ class TokenizerManager:
             trust_remote_code=server_args.trust_remote_code,
             model_overide_args=model_overide_args,
         )
+
         if server_args.context_length is not None:
             self.context_len = server_args.context_length
         else:
@@ -133,50 +149,54 @@ class TokenizerManager:
             async for response in self._handle_batch_request(obj, request):
                 yield response
 
-    async def _handle_single_request(self, obj, request, index=None, is_prefill=False):
-        if is_prefill:
-            if isinstance(obj.text, list):
-                input_text = obj.text[index]
-                rid = obj.rid[index]
-            else:
-                input_text = obj.text
-                rid = obj.rid[0]
-            input_ids = self.tokenizer.encode(input_text)
-            sampling_params = SamplingParams(**obj.sampling_params[0])
-            sampling_params.max_new_tokens = 0
-            pixel_values, image_hash, image_size = await self._get_pixel_values(
-                obj.image_data[0]
-            )
-            return_logprob = obj.return_logprob[0]
-            logprob_start_len = obj.logprob_start_len[0]
-            top_logprobs_num = obj.top_logprobs_num[0]
-        else:
-            rid = obj.rid if index is None else obj.rid[index]
-            input_text = obj.text if index is None else obj.text[index]
+    async def _handle_single_request(
+        self, obj, request, index=None, is_cache_for_prefill=False
+    ):
+        if not is_cache_for_prefill:
+            not_use_index = not (index is not None)
+            rid = obj.rid if not_use_index else obj.rid[index]
+            input_text = obj.text if not_use_index else obj.text[index]
             input_ids = (
                 self.tokenizer.encode(input_text)
                 if obj.input_ids is None
                 else obj.input_ids
             )
-            if index is not None and obj.input_ids:
+            if not not_use_index and obj.input_ids:
                 input_ids = obj.input_ids[index]
 
             self._validate_input_length(input_ids)
+
             sampling_params = self._get_sampling_params(
-                obj.sampling_params if index is None else obj.sampling_params[index]
+                obj.sampling_params if not_use_index else obj.sampling_params[index]
             )
             pixel_values, image_hash, image_size = await self._get_pixel_values(
-                obj.image_data if index is None else obj.image_data[index]
+                obj.image_data if not_use_index else obj.image_data[index]
             )
             return_logprob = (
-                obj.return_logprob if index is None else obj.return_logprob[index]
+                obj.return_logprob if not_use_index else obj.return_logprob[index]
             )
             logprob_start_len = (
-                obj.logprob_start_len if index is None else obj.logprob_start_len[index]
+                obj.logprob_start_len if not_use_index else obj.logprob_start_len[index]
             )
             top_logprobs_num = (
-                obj.top_logprobs_num if index is None else obj.top_logprobs_num[index]
+                obj.top_logprobs_num if not_use_index else obj.top_logprobs_num[index]
             )
+        else:
+            if isinstance(obj.text, list):
+                input_text = obj.text[index]
+                rid = obj.rid[index]
+            else:
+                input_text = obj.text
+                rid = obj.rid[0]
+            input_ids = self.tokenizer.encode(input_text)
+            sampling_params = SamplingParams(**obj.sampling_params[0])
+            sampling_params.max_new_tokens = 0
+            pixel_values, image_hash, image_size = await self._get_pixel_values(
+                obj.image_data[0]
+            )
+            return_logprob = obj.return_logprob[0]
+            logprob_start_len = obj.logprob_start_len[0]
+            top_logprobs_num = obj.top_logprobs_num[0]
 
         tokenized_obj = TokenizedGenerateReqInput(
             rid,
@@ -196,26 +216,26 @@ class TokenizerManager:
         event = asyncio.Event()
         state = ReqState([], False, event)
         self.rid_to_state[rid] = state
-        if is_prefill:
-            await self._wait_for_prefill_response(event, state, obj, request, rid)
-            yield input_ids
-        else:
+        if not is_cache_for_prefill:
             async for response in self._wait_for_response(
                 event, state, obj, rid, request
             ):
                 yield response
+        else:
+            await self._wait_for_cache_prefill_response(event, state, obj, rid, request)
+            yield input_ids
 
-    async def _handle_batch_request(self, obj, request):
+    async def _handle_batch_request(self, obj: GenerateReqInput, request):
         batch_size = obj.batch_size
-        parallel_sample_num = obj.sampling_params[0].get("n", 1)
+        parallel_sample_num = obj.parallel_sample_num
 
         if parallel_sample_num != 1:
-            ## send prefill requests
+            # Send prefill requests to cache the common input
             parallel_sample_num += 1
             input_id_result = [] if obj.input_ids is None else None
             for i in range(batch_size):
                 async for input_id in self._handle_single_request(
-                    obj, request, index=i, is_prefill=True
+                    obj, request, index=i, is_cache_for_prefill=True
                 ):
                     if input_id_result is not None:
                         input_id_result.append(input_id)
@@ -231,7 +251,7 @@ class TokenizerManager:
                     continue
                 index = i * parallel_sample_num + j
                 if parallel_sample_num != 1:
-                    # Here when using parallel sampling we shoul consider prefill stage so the index is :  j + i * (parallel_sample_num-1) + batch_size - 1
+                    # Here when using parallel sampling we should consider prefill stage so the index is :  j + i * (parallel_sample_num-1) + batch_size - 1
                     index += batch_size - 1 - i
                 rid = obj.rid[index]
                 if parallel_sample_num == 1:
@@ -308,17 +328,15 @@ class TokenizerManager:
 
         yield output_list
 
-    def _validate_input_length(self, input_ids):
+    def _validate_input_length(self, input_ids: List[int]):
         if len(input_ids) >= self.context_len:
             raise ValueError(
                 f"The input ({len(input_ids)} tokens) is longer than the "
                 f"model's context length ({self.context_len} tokens)."
             )
 
-    def _get_sampling_params(self, sampling_params_data, max_new_tokens=None):
+    def _get_sampling_params(self, sampling_params_data: dict):
         sampling_params = SamplingParams(**sampling_params_data)
-        if max_new_tokens is not None:
-            sampling_params.max_new_tokens = max_new_tokens
         if sampling_params.max_new_tokens != 0:
             sampling_params.normalize(self.tokenizer)
             sampling_params.verify()
@@ -332,7 +350,14 @@ class TokenizerManager:
         else:
             return None, None, None
 
-    async def _wait_for_response(self, event, state, obj, rid, request):
+    async def _wait_for_response(
+        self,
+        event: asyncio.Event,
+        state: ReqState,
+        obj: GenerateReqInput,
+        rid: str,
+        request,
+    ):
         while True:
             try:
                 await asyncio.wait_for(event.wait(), timeout=4)
@@ -361,7 +386,14 @@ class TokenizerManager:
             event.clear()
             yield out
 
-    async def _wait_for_prefill_response(self, event, state, obj, request, rid):
+    async def _wait_for_cache_prefill_response(
+        self,
+        event: asyncio.Event,
+        state: ReqState,
+        obj: GenerateReqInput,
+        rid: str,
+        request,
+    ):
         while True:
             try:
                 await asyncio.wait_for(state.event.wait(), timeout=4)
@@ -380,7 +412,7 @@ class TokenizerManager:
         req = FlushCacheReq()
         self.send_to_router.send_pyobj(req)
 
-    def abort_request(self, rid):
+    def abort_request(self, rid: str):
         if rid not in self.rid_to_state:
             return
         del self.rid_to_state[rid]
@@ -426,31 +458,37 @@ class TokenizerManager:
                 state.event.set()
 
     def convert_logprob_style(
-        self, ret, return_logprob, top_logprobs_num, return_text_in_logprobs
+        self,
+        ret: dict,
+        return_logprob: bool,
+        top_logprobs_num: int,
+        return_text_in_logprobs: bool,
     ):
         if return_logprob:
-            ret["meta_info"]["prefill_token_logprobs"] = self.detokenize_logprob_tokens(
-                ret["meta_info"]["prefill_token_logprobs"], return_text_in_logprobs
+            ret["meta_info"]["input_token_logprobs"] = self.detokenize_logprob_tokens(
+                ret["meta_info"]["input_token_logprobs"], return_text_in_logprobs
             )
-            ret["meta_info"]["decode_token_logprobs"] = self.detokenize_logprob_tokens(
-                ret["meta_info"]["decode_token_logprobs"], return_text_in_logprobs
+            ret["meta_info"]["output_token_logprobs"] = self.detokenize_logprob_tokens(
+                ret["meta_info"]["output_token_logprobs"], return_text_in_logprobs
             )
 
             if top_logprobs_num > 0:
-                ret["meta_info"]["prefill_top_logprobs"] = (
+                ret["meta_info"]["input_top_logprobs"] = (
                     self.detokenize_top_logprobs_tokens(
-                        ret["meta_info"]["prefill_top_logprobs"],
+                        ret["meta_info"]["input_top_logprobs"],
                         return_text_in_logprobs,
                     )
                 )
-                ret["meta_info"]["decode_top_logprobs"] = (
+                ret["meta_info"]["output_top_logprobs"] = (
                     self.detokenize_top_logprobs_tokens(
-                        ret["meta_info"]["decode_top_logprobs"], return_text_in_logprobs
+                        ret["meta_info"]["output_top_logprobs"], return_text_in_logprobs
                     )
                 )
         return ret
 
-    def detokenize_logprob_tokens(self, token_logprobs, decode_to_text):
+    def detokenize_logprob_tokens(
+        self, token_logprobs: List[Tuple[float, int]], decode_to_text: bool
+    ):
         if not decode_to_text:
             return [(logprob, token_id, None) for logprob, token_id in token_logprobs]
 
@@ -461,10 +499,14 @@ class TokenizerManager:
             for (logprob, token_id), token_text, in zip(token_logprobs, token_texts)
         ]
 
-    def detokenize_top_logprobs_tokens(self, top_logprobs, decode_to_text):
-        for i, t in enumerate(top_logprobs):
-            if t:
-                top_logprobs[i] = self.detokenize_logprob_tokens(t, decode_to_text)
+    def detokenize_top_logprobs_tokens(self, top_logprobs, decode_to_text: bool):
+        # TODO: The current implementation only batches the detokenization for top-k tokens per single position.
+        # We should batch all top-k tokens in all positions.
+        for i, token_top_logprobs in enumerate(top_logprobs):
+            if token_top_logprobs:
+                top_logprobs[i] = self.detokenize_logprob_tokens(
+                    token_top_logprobs, decode_to_text
+                )
         return top_logprobs