Add max_prefill_num_token into server arguments (#133)

2024-02-03 02:35:54 -08:00
parent 67be11c790
commit e095b16236
3 changed files with 12 additions and 2 deletions
@@ -82,7 +82,8 @@ class ModelRpcServer(rpyc.Service):
        self.max_total_num_token = self.model_runner.max_total_num_token
        self.max_num_running_seq = self.max_total_num_token // 2
        self.max_prefill_num_token = max(
-            self.model_config.context_len, self.max_total_num_token // 6
+            self.model_config.context_len,
+            self.max_total_num_token // 6 if server_args.max_prefill_num_token is None else server_args.max_prefill_num_token,
        )
        self.int_token_logit_bias = torch.tensor(
            get_int_token_logit_bias(self.tokenizer, self.model_config.vocab_size)