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

sglang/lang/backend/runtime_endpoint.py

@@ -253,14 +253,14 @@ class RuntimeEndpoint(BaseBackend):
             r["meta_info"]["normalized_prompt_logprob"] for r in obj
         ]
         decision = choices[np.argmax(normalized_prompt_logprobs)]
-        prefill_token_logprobs = [r["meta_info"]["prefill_token_logprobs"] for r in obj]
-        decode_token_logprobs = [r["meta_info"]["decode_token_logprobs"] for r in obj]
+        input_token_logprobs = [r["meta_info"]["input_token_logprobs"] for r in obj]
+        output_token_logprobs = [r["meta_info"]["output_token_logprobs"] for r in obj]
 
         return (
             decision,
             normalized_prompt_logprobs,
-            prefill_token_logprobs,
-            decode_token_logprobs,
+            input_token_logprobs,
+            output_token_logprobs,
         )
 
     def concatenate_and_append(self, src_rids: List[str], dst_rid: str):

sglang/lang/interpreter.py

@@ -541,16 +541,16 @@ class StreamExecutor:
         (
             decision,
             normalized_prompt_logprobs,
-            prefill_token_logprobs,
-            decode_token_logprobs,
+            input_token_logprobs,
+            output_token_logprobs,
         ) = self.backend.select(self, expr.choices, expr.temperature)
         if expr.name is not None:
             name = expr.name
             self.variables[name] = decision
             self.meta_info[name] = {
                 "normalized_prompt_logprobs": normalized_prompt_logprobs,
-                "prefill_token_logprobs": prefill_token_logprobs,
-                "decode_token_logprobs": decode_token_logprobs,
+                "input_token_logprobs": input_token_logprobs,
+                "output_token_logprobs": output_token_logprobs,
             }
             self.variable_event[name].set()
         self.text_ += decision

sglang/srt/constrained/fsm_cache.py

@@ -21,7 +21,27 @@ class FSMCache(BaseCache):
             tokenizer = AutoTokenizer.from_pretrained(
                 tokenizer_path, **tokenizer_args_dict
             )
-            self.outlines_tokenizer = TransformerTokenizer(tokenizer)
+            try:
+                self.outlines_tokenizer = TransformerTokenizer(tokenizer)
+            except AttributeError:
+                # FIXME: tmp fix for chatglm2 & chatglm3 (pad_token_id=0)
+                origin_pad_token_id = tokenizer.pad_token_id
+
+                def fset(self, value):
+                    self._value = value
+
+                type(tokenizer).pad_token_id = property(
+                    fget=type(tokenizer).pad_token_id.fget, fset=fset
+                )
+                self.outlines_tokenizer = TransformerTokenizer(tokenizer)
+                self.outlines_tokenizer.tokenizer.pad_token_id = origin_pad_token_id
+                self.outlines_tokenizer.pad_token_id = origin_pad_token_id
+                self.outlines_tokenizer.pad_token = (
+                    self.outlines_tokenizer.tokenizer.pad_token
+                )
+                self.outlines_tokenizer.vocabulary = (
+                    self.outlines_tokenizer.tokenizer.get_vocab()
+                )
         else:
             self.outlines_tokenizer = TransformerTokenizer(
                 tokenizer_path, **tokenizer_args_dict

sglang/srt/hf_transformers_utils.py

@@ -73,7 +73,9 @@ def get_context_length(config):
     rope_scaling = getattr(config, "rope_scaling", None)
     if rope_scaling:
         rope_scaling_factor = config.rope_scaling["factor"]
-        if config.rope_scaling["rope_type"] == "llama3":
+        if "original_max_position_embeddings" in rope_scaling:
+            rope_scaling_factor = 1
+        if config.rope_scaling.get("rope_type", None) == "llama3":
             rope_scaling_factor = 1
     else:
         rope_scaling_factor = 1

sglang/srt/layers/logits_processor.py

@@ -1,7 +1,7 @@
 """Logits processing."""
 
 import dataclasses
-from typing import List, Union
+from typing import List, Optional, Union
 
 import torch
 from torch import nn
@@ -22,23 +22,23 @@ class LogitProcessorOutput:
 
     # 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 logprobs of input tokens.      shape: [#token, vocab_size]
+    input_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
+    # The logprob and id of the top-k tokens in input positions.  shape [#seq, #token, k] of Tuple(logprob, token_id)
+    input_top_logprobs: List
+    # The logprob and id of the top-k tokens in output positions. shape [#seq, #token, k] of Tuple(logprob, token_id)
+    output_top_logprobs: List
 
 
 @dataclasses.dataclass
 class LogitsMetadata:
     forward_mode: ForwardMode
-    return_logprob: bool
+    return_logprob: bool = False
 
-    extend_seq_lens: torch.Tensor = None
-    extend_start_loc: torch.Tensor = None
-    top_logprobs_nums: List[int] = None
+    extend_seq_lens: Optional[torch.Tensor] = None
+    extend_start_loc: Optional[torch.Tensor] = None
+    top_logprobs_nums: Optional[List[int]] = None
 
     @classmethod
     def from_input_metadata(cls, input_metadata: InputMetadata):
@@ -58,20 +58,16 @@ class LogitsProcessor(nn.Module):
         self.tp_size = get_tensor_model_parallel_world_size()
 
     def _get_normalized_prompt_logprobs(
-        self, prefill_token_logprobs, logits_metadata: LogitsMetadata
+        self, input_token_logprobs, logits_metadata: LogitsMetadata
     ):
-        logprobs_cumsum = torch.cumsum(
-            prefill_token_logprobs, dim=0, dtype=torch.float32
-        )
+        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.clamp_(min=0, max=prefill_token_logprobs.shape[0] - 1)
-        end.clamp_(min=0, max=prefill_token_logprobs.shape[0] - 1)
+        start.clamp_(min=0, max=input_token_logprobs.shape[0] - 1)
+        end.clamp_(min=0, max=input_token_logprobs.shape[0] - 1)
         sum_logp = (
-            logprobs_cumsum[end]
-            - logprobs_cumsum[start]
-            + prefill_token_logprobs[start]
+            logprobs_cumsum[end] - logprobs_cumsum[start] + input_token_logprobs[start]
         )
         normalized_prompt_logprobs = sum_logp / (
             (logits_metadata.extend_seq_lens - 1).clamp(min=1)
@@ -79,37 +75,38 @@ class LogitsProcessor(nn.Module):
 
         return normalized_prompt_logprobs
 
-    def _get_top_logprobs(self, all_logprobs, logits_metadata: LogitsMetadata):
+    @staticmethod
+    def get_top_logprobs(all_logprobs, logits_metadata: LogitsMetadata):
         # TODO: vectorize the code below
         if logits_metadata.forward_mode == ForwardMode.DECODE:
-            decode_top_logprobs = []
+            output_top_logprobs = []
             for i in range(all_logprobs.shape[0]):
                 k = logits_metadata.top_logprobs_nums[i]
                 t = all_logprobs[i].topk(k)
                 v_cpu = t.values.tolist()
                 p_cpu = t.indices.tolist()
-                decode_top_logprobs.append(list(zip(v_cpu, p_cpu)))
-            return None, decode_top_logprobs
+                output_top_logprobs.append(list(zip(v_cpu, p_cpu)))
+            return None, output_top_logprobs
         else:
-            prefill_top_logprobs, decode_top_logprobs = [], []
+            input_top_logprobs, output_top_logprobs = [], []
             pt = 0
             extend_seq_lens_cpu = logits_metadata.extend_seq_lens.tolist()
             for i, extend_seq_len in enumerate(extend_seq_lens_cpu):
                 if extend_seq_len == 0:
-                    prefill_top_logprobs.append([])
-                    decode_top_logprobs.append([])
+                    input_top_logprobs.append([])
+                    output_top_logprobs.append([])
                     continue
                 k = logits_metadata.top_logprobs_nums[i]
                 t = all_logprobs[pt : pt + extend_seq_len].topk(k)
                 vs_cpu = t.values.tolist()
                 ps_cpu = t.indices.tolist()
-                prefill_top_logprobs.append(
+                input_top_logprobs.append(
                     [list(zip(vs_cpu[j], ps_cpu[j])) for j in range(len(vs_cpu) - 1)]
                 )
-                decode_top_logprobs.append(list(zip(vs_cpu[-1], ps_cpu[-1])))
+                output_top_logprobs.append(list(zip(vs_cpu[-1], ps_cpu[-1])))
                 pt += extend_seq_len
 
-            return prefill_top_logprobs, decode_top_logprobs
+            return input_top_logprobs, output_top_logprobs
 
     def forward(
         self,
@@ -136,7 +133,7 @@ class LogitsProcessor(nn.Module):
         last_logits = torch.matmul(last_hidden, weight.T)
         if self.tp_size > 1:
             last_logits = tensor_model_parallel_all_gather(last_logits)
-        last_logits = last_logits[:, : self.config.vocab_size]
+        last_logits = last_logits[:, : self.config.vocab_size].float()
 
         if hasattr(self.config, "final_logit_softcapping"):
             last_logits /= self.config.final_logit_softcapping
@@ -149,63 +146,75 @@ class LogitsProcessor(nn.Module):
                 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,
+                input_token_logprobs=None,
+                input_top_logprobs=None,
+                output_top_logprobs=None,
             )
         else:
             # When logprob is requested, compute the logits for all tokens.
             if logits_metadata.forward_mode == ForwardMode.DECODE:
-                all_logits = last_logits
-            else:
-                all_logits = torch.matmul(hidden_states, weight.T)
-                if self.tp_size > 1:
-                    all_logits = tensor_model_parallel_all_gather(all_logits)
-                all_logits = all_logits[:, : self.config.vocab_size]
-
-            all_logprobs = all_logits.float()
-            del all_logits
-            all_logprobs[:] = torch.nn.functional.log_softmax(all_logprobs, dim=-1)
+                last_logprobs = torch.nn.functional.log_softmax(last_logits, dim=-1)
 
-            # Get the logprob of top-k tokens
-            return_top_logprob = any(x > 0 for x in logits_metadata.top_logprobs_nums)
-            if return_top_logprob:
-                prefill_top_logprobs, decode_top_logprobs = self._get_top_logprobs(
-                    all_logprobs, logits_metadata
+                # Get the logprob of top-k tokens
+                return_top_logprob = any(
+                    x > 0 for x in logits_metadata.top_logprobs_nums
                 )
-            else:
-                prefill_top_logprobs = decode_top_logprobs = None
+                if return_top_logprob:
+                    output_top_logprobs = self.get_top_logprobs(
+                        last_logprobs, logits_metadata
+                    )[1]
+                else:
+                    output_top_logprobs = None
 
-            if logits_metadata.forward_mode == ForwardMode.DECODE:
                 return LogitProcessorOutput(
                     next_token_logits=last_logits,
-                    next_token_logprobs=all_logprobs,
+                    next_token_logprobs=last_logprobs,
                     normalized_prompt_logprobs=None,
-                    prefill_token_logprobs=None,
-                    prefill_top_logprobs=None,
-                    decode_top_logprobs=decode_top_logprobs,
+                    input_token_logprobs=None,
+                    input_top_logprobs=None,
+                    output_top_logprobs=output_top_logprobs,
                 )
             else:
+                all_logits = torch.matmul(hidden_states, weight.T)
+                if self.tp_size > 1:
+                    all_logits = tensor_model_parallel_all_gather(all_logits)
+                all_logits = all_logits[:, : self.config.vocab_size].float()
+
+                all_logprobs = all_logits
+                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 logits_metadata.top_logprobs_nums
+                )
+                if return_top_logprob:
+                    input_top_logprobs, output_top_logprobs = self.get_top_logprobs(
+                        all_logprobs, logits_metadata
+                    )
+                else:
+                    input_top_logprobs = output_top_logprobs = None
+
                 last_logprobs = all_logprobs[last_index]
 
                 # 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.
-                prefill_token_logprobs = all_logprobs[
+                input_token_logprobs = all_logprobs[
                     torch.arange(all_logprobs.shape[0], device="cuda"),
                     torch.cat([input_ids[1:], torch.tensor([0], device="cuda")]),
                 ]
 
                 normalized_prompt_logprobs = self._get_normalized_prompt_logprobs(
-                    prefill_token_logprobs, logits_metadata
+                    input_token_logprobs, logits_metadata
                 )
 
                 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,
+                    input_token_logprobs=input_token_logprobs,
+                    input_top_logprobs=input_top_logprobs,
+                    output_top_logprobs=output_top_logprobs,
                 )
 
 

sglang/srt/layers/radix_attention.py

@@ -7,8 +7,11 @@ from torch import nn
 from sglang.global_config import global_config
 from sglang.srt.layers.extend_attention import extend_attention_fwd
 from sglang.srt.layers.token_attention import token_attention_fwd
-from sglang.srt.managers.controller.model_runner import ForwardMode, InputMetadata
-from sglang.srt.server import global_server_args_dict
+from sglang.srt.managers.controller.model_runner import (
+    ForwardMode,
+    InputMetadata,
+    global_server_args_dict,
+)
 
 
 class RadixAttention(nn.Module):

sglang/srt/layers/token_attention.py

@@ -5,7 +5,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.server import global_server_args_dict
+from sglang.srt.managers.controller.infer_batch import global_server_args_dict
 
 if global_server_args_dict.get("attention_reduce_in_fp32", False):
     REDUCE_TRITON_TYPE = tl.float32

sglang/srt/managers/controller/cuda_graph_runner.py

@@ -9,7 +9,11 @@ from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
 from vllm.distributed.parallel_state import graph_capture
 from vllm.model_executor.custom_op import CustomOp
 
-from sglang.srt.layers.logits_processor import LogitProcessorOutput
+from sglang.srt.layers.logits_processor import (
+    LogitProcessorOutput,
+    LogitsMetadata,
+    LogitsProcessor,
+)
 from sglang.srt.managers.controller.infer_batch import (
     Batch,
     ForwardMode,
@@ -185,7 +189,6 @@ class CudaGraphRunner:
 
     def replay(self, batch: Batch):
         assert batch.out_cache_loc is not None
-        assert not batch.return_logprob
         raw_bs = len(batch.reqs)
 
         # Pad
@@ -218,23 +221,29 @@ class CudaGraphRunner:
         output = self.output_buffers[bs]
 
         # Unpad
-        if bs == raw_bs:
-            return output
-        else:
+        if bs != raw_bs:
             output = LogitProcessorOutput(
                 next_token_logits=output.next_token_logits[:raw_bs],
-                next_token_logprobs=(
-                    output.next_token_logprobs[:raw_bs]
-                    if output.next_token_logprobs is not None
-                    else None
-                ),
+                next_token_logprobs=None,
                 normalized_prompt_logprobs=None,
-                prefill_token_logprobs=None,
-                prefill_top_logprobs=None,
-                decode_top_logprobs=(
-                    output.decode_top_logprobs[:raw_bs]
-                    if output.decode_top_logprobs is not None
-                    else None
-                ),
+                input_token_logprobs=None,
+                input_top_logprobs=None,
+                output_top_logprobs=None,
             )
+
+        # Extract logprobs
+        if batch.return_logprob:
+            output.next_token_logprobs = torch.nn.functional.log_softmax(
+                output.next_token_logits, dim=-1
+            )
+            return_top_logprob = any(x > 0 for x in batch.top_logprobs_nums)
+            if return_top_logprob:
+                logits_metadata = LogitsMetadata(
+                    forward_mode=ForwardMode.DECODE,
+                    top_logprobs_nums=batch.top_logprobs_nums,
+                )
+                output.output_top_logprobs = LogitsProcessor.get_top_logprobs(
+                    output.next_token_logprobs, logits_metadata
+                )[1]
+
         return output

sglang/srt/managers/controller/infer_batch.py

@@ -17,6 +17,13 @@ from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
 
 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,
+}
+
 
 class ForwardMode(IntEnum):
     # Prefill a new sequence. This is deprecated now. "EXTEND" covers this case.
@@ -124,10 +131,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 +251,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
 
@@ -376,7 +383,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(
@@ -687,13 +694,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)
@@ -933,3 +948,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/controller/model_runner.py

@@ -25,7 +25,12 @@ from vllm.distributed import (
 from vllm.model_executor.models import ModelRegistry
 
 from sglang.global_config import global_config
-from sglang.srt.managers.controller.infer_batch import Batch, ForwardMode, InputMetadata
+from sglang.srt.managers.controller.infer_batch import (
+    Batch,
+    ForwardMode,
+    InputMetadata,
+    global_server_args_dict,
+)
 from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
@@ -60,7 +65,13 @@ class ModelRunner:
         self.nccl_port = nccl_port
         self.server_args = server_args
         self.is_multimodal_model = is_multimodal_model(self.model_config)
-        monkey_patch_vllm_dummy_weight_loader()
+        global_server_args_dict.update(
+            {
+                "disable_flashinfer": server_args.disable_flashinfer,
+                "disable_flashinfer_sampling": server_args.disable_flashinfer_sampling,
+                "attention_reduce_in_fp32": server_args.attention_reduce_in_fp32,
+            }
+        )
 
         # Init torch distributed
         torch.cuda.set_device(self.gpu_id)
@@ -95,7 +106,7 @@ class ModelRunner:
 
         # Load the model and create memory pool
         self.load_model()
-        self.init_memory_pool(total_gpu_memory)
+        self.init_memory_pool(total_gpu_memory, server_args.max_num_reqs)
         self.init_cublas()
         self.init_flash_infer()
 
@@ -108,6 +119,7 @@ class ModelRunner:
             f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
         )
 
+        monkey_patch_vllm_dummy_weight_loader()
         device_config = DeviceConfig()
         load_config = LoadConfig(load_format=self.server_args.load_format)
         vllm_model_config = VllmModelConfig(
@@ -176,7 +188,7 @@ class ModelRunner:
         max_num_token = int(rest_memory * (1 << 30) // cell_size)
         return max_num_token
 
-    def init_memory_pool(self, total_gpu_memory):
+    def init_memory_pool(self, total_gpu_memory, max_num_reqs=None):
         self.max_total_num_tokens = self.profile_max_num_token(total_gpu_memory)
 
         if self.max_total_num_tokens <= 0:
@@ -184,11 +196,14 @@ class ModelRunner:
                 "Not enough memory. Please try to increase --mem-fraction-static."
             )
 
-        self.req_to_token_pool = ReqToTokenPool(
-            max(
+        if max_num_reqs is None:
+            max_num_reqs = max(
                 int(self.max_total_num_tokens / self.model_config.context_len * 512),
                 2048,
-            ),
+            )
+
+        self.req_to_token_pool = ReqToTokenPool(
+            max_num_reqs,
             self.model_config.context_len + 8,
         )
         self.token_to_kv_pool = TokenToKVPool(