Truely use maxd; Optimize query start loc; Make evaluation use LLM (#86)

* Truely use maxd in Token Throttling

* Optimize query_start_loc

* Make evaluation use offline LLM

Author: gty111

README.md

@@ -137,8 +137,7 @@ python benchmarks/benchmark_prefix_serving.py \
 
 ### Evaluate Output Quality
 ```
-# Launch server first
-python evaluations/evaluate_MMLU_pro.py --model $MODEL --port $PORT
+python evaluations/evaluate_MMLU_pro.py --model $MODEL
 ```
 
 ## Supported Models

evaluations/evaluate_MMLU_pro.py benchmarks/evaluate_MMLU_pro.pyevaluations/evaluate_MMLU_pro.py renamed to benchmarks/evaluate_MMLU_pro.py

@@ -1,18 +1,15 @@
 # adopt from https://github.com/TIGER-AI-Lab/MMLU-Pro/blob/main/evaluate_from_api.py
-
-import os
 import re
 import random
-from tqdm import tqdm
-from datasets import load_dataset
 import argparse
-from benchmarks.backend_request_func import async_request_openai_chat_completions, RequestFuncInput
 import asyncio
 
+from gllm import LLM
 
-API_KEY = "EMPTY"
-random.seed(12345)
+from tqdm import tqdm
+from datasets import load_dataset
 
+random.seed(12345)
 
 def load_mmlu_pro():
     dataset = load_dataset("TIGER-Lab/MMLU-Pro")
@@ -83,7 +80,7 @@ def extract_final(text):
         return None
 
 
-def single_request(api_url, single_question, cot_examples_dict, pbar):
+def single_request(single_question, cot_examples_dict):
     category = single_question["category"]
     cot_examples = cot_examples_dict[category]
     question = single_question["question"]
@@ -95,43 +92,44 @@ def single_request(api_url, single_question, cot_examples_dict, pbar):
         prompt += format_example(each["question"],
                                  each["options"], each["cot_content"])
     input_text = format_example(question, options)
-
     prompt = prompt + input_text
-
-    request_func_input = RequestFuncInput(prompt=prompt,
-                                          api_url=api_url,
-                                          prompt_len=len(prompt),
-                                          output_len=args.output_len,
-                                          model=args.model,
-                                          )
-    return async_request_openai_chat_completions(request_func_input=request_func_input, pbar=pbar)
+    
+    return prompt
 
 
 
 async def evaluate(subjects):
-    api_url = f"http://{args.host}:{args.port}/v1/chat/completions"
     test_df, dev_df = load_mmlu_pro()
     if not subjects:
         subjects = list(test_df.keys())
     print("assigned subjects", subjects)
     category_record = {'total':{'#correct':0,'#wrong':0}}
 
-    print(f"Sending requests ...")
-    pbar = tqdm()
-    tasks = []
+    llm = LLM(model_path=args.model,
+              gpu_memory_util=args.gpu_memory_util,
+              kvthresh=args.kvthresh,
+              pp_size=args.pp,
+              tp_size=args.tp,
+              enable_prefix_caching=True,
+              use_thinking=False)
+    
+    print(f"generating requests ...")
+    prompts = []
     test_data_total = []
     for subject in subjects:
         test_data = test_df[subject][:args.num_per_sub]
         test_data_total.extend(test_data)
         for each in test_data:
-            tasks.append(single_request(api_url, each, dev_df, pbar))
-    pbar.total = len(tasks)
-    completions = await asyncio.gather(*tasks)
-    pbar.close()
+            prompts.append(single_request(each, dev_df))
+    
+    seqs = llm.generate(prompts, output_lens=[args.output_len for i in range(len(prompts))])
+    
+    outputs = [seq.output for seq in seqs]
+    
     print(f"Processing completions ...")
-    for idx, each in tqdm(enumerate(test_data_total),total=len(tasks)):
+    for idx, each in tqdm(enumerate(test_data_total),total=len(prompts)):
         label = each["answer"]
-        response = completions[idx].generated_text
+        response = outputs[idx]
         response = response.replace('**', '')
         pred = extract_answer(response)
         category = each["category"]
@@ -162,8 +160,10 @@ async def evaluate(subjects):
     parser.add_argument("--assigned_subjects", "-a", type=str, default="all",
                         help="business, law, psychology, biology, chemistry, history, other, health, "
                              "economics, math, physics, computer science, philosophy, engineering")
-    parser.add_argument("--host", type=str, default='0.0.0.0')
-    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--tp", type=int, default=1)
+    parser.add_argument("--pp", type=int, default=1)
+    parser.add_argument('--gpu-memory-util', type=float, default=0.9)
+    parser.add_argument('--kvthresh', type=float, default=0.2)
     parser.add_argument("--output-len", type=int, default=1024)
     parser.add_argument("--num-per-sub", type=int, default=100)
     assigned_subjects = []

gllm/entrypoints/api_server.py

@@ -102,11 +102,11 @@ async def run_server(args):
     parser.add_argument('--disable-ep', help='Disable expert parallelism (EP is enable by default)', action='store_true')
     parser.add_argument('--assigned-layers', type=str, help='If the model have 64 layers, we can set it to 16,16,16,16 or 16,16,17,15', default=None)
     # Token Throttling
-    parser.add_argument('--maxd', type=int, help='Maximum decode token count, used in LLM (offline infernce)', default=512)
-    parser.add_argument('--maxp', type=int, help='Maximum token count in prefill', default=2048)
-    parser.add_argument('--minp', type=int, help='Minimum token count in prefill, used in PipeAsyncLLM', default=32)
-    parser.add_argument('--iterp', type=int, help='Number of iterations to process waiting prefill tokens, used in PipeAsyncLLM', default=8)
-    parser.add_argument('--kvthresh', type=float, help='KV cache threshold for prefill operations', default=0.05)
+    parser.add_argument('--maxd', type=int, help='Maximum decode token count per batch (Token Throttling)', default=2048)
+    parser.add_argument('--maxp', type=int, help='Maximum prefill token count per batch (Token Throttling) or token budget in Sarathi-Serve', default=2048)
+    parser.add_argument('--minp', type=int, help='Minimum prefill token count per batch (Token Throttling)', default=32)
+    parser.add_argument('--iterp', type=int, help='Number of iterations to process waiting prefill tokens (Token Throttling)', default=8)
+    parser.add_argument('--kvthresh', type=float, help='KV cache threshold for prefill operations (Token Throttling)', default=0.05)
     parser.add_argument('--use-naive-schedule', help='Use scheduling policy in Sarathi-Serve', action='store_true')
     # Multi-Node deployment
     parser.add_argument('--launch-mode', type=str, choices=['normal', 'master', 'slave'], default='normal')

gllm/input_data.py

@@ -29,9 +29,9 @@ def __init__(self, seqs: List[Sequence], memory_manager: MemoryManager):
         self.slot_mapping_tensor = self.get_slot_mapping()
         self.tokens = self.get_tokens()
         self.positions = self.get_position()
-        self.max_seq_len, self.seq_start_loc = self.get_seq_len_loc()
+        self.max_seq_len, self.seq_start_loc = self.get_seq_lens()
         self.block_table = self.get_block_table()
-        self.max_query_len, self.query_start_loc = self.get_query_len_loc()
+        self.max_query_len, self.query_start_loc = self.get_query_start_loc()
 
         assert self.tokens.shape == self.positions.shape
 
@@ -50,21 +50,18 @@ def get_position(self):
         return async_tensor_h2d(
             positions_list, torch.long, 'cuda', True)
 
-    def get_seq_len_loc(self):
-        seq_start_loc = [seq.seq_len for seq in self.seqs]
-        max_seqlen = max(seq_start_loc)
-        return max_seqlen, async_tensor_h2d(seq_start_loc, torch.int32, 'cuda', True)
+    def get_seq_lens(self):
+        seq_lens = [seq.seq_len for seq in self.seqs]
+        max_seqlen = max(seq_lens)
+        return max_seqlen, async_tensor_h2d(seq_lens, torch.int32, 'cuda', True)
 
-    def get_query_len_loc(self):
-        max_query_len = 0
-        cu_query_len = 0
-        query_start_loc = [0]
-        for seq in self.seqs:
-            query_len = seq.to_compute_token_num
-            cu_query_len += query_len
-            query_start_loc.append(cu_query_len)
-            max_query_len = max(query_len, max_query_len)
-        return max_query_len, async_tensor_h2d(query_start_loc, torch.int32, 'cuda', True)
+    def get_query_start_loc(self):
+        query_lens = [0] + [seq.to_compute_token_num for seq in self.seqs]
+        max_query_len = max(query_lens)
+        query_start_loc = torch.from_numpy(np.cumsum(query_lens)).to(device='cuda', 
+                                                                    dtype=torch.int32, 
+                                                                    non_blocking=True)
+        return max_query_len, query_start_loc
 
     def get_block_table(self):
         block_tables_list = [seq.page_table for seq in self.seqs]
@@ -73,7 +70,7 @@ def get_block_table(self):
             (len(block_tables_list), max_num_block), 0, dtype=np.int32)
         for idx, block_table in enumerate(block_tables_list):
             block_tables[idx, :len(block_table)] = block_table
-        return torch.from_numpy(block_tables).to(device='cuda',non_blocking=True)
+        return torch.from_numpy(block_tables).to(device='cuda', non_blocking=True)
 
     def get_slot_mapping(self):
         slot_mapping = []

gllm/worker.py

@@ -65,6 +65,7 @@ def init(self):
                 self.pp_size,
                 self.model_runner.memory_manager, 
                 self.use_naive_schedule,
+                self.model_runner.maxd,
                 self.model_runner.maxp,
                 self.model_runner.minp,
                 self.model_runner.iterp,

gllm/worker_scheduler.py

@@ -13,10 +13,12 @@
 
 
 class WorkerScheduler():
-    def __init__(self, pp_size, memory_manager:MemoryManager, use_naive_schedule, maxp, minp, iterp, page_size, kvthresh):
+    def __init__(self, pp_size, memory_manager:MemoryManager, use_naive_schedule, 
+                 maxd, maxp, minp, iterp, page_size, kvthresh):
         self.pp_size = pp_size
         self.memory_manager = memory_manager
         self.use_naive_schedule = use_naive_schedule
+        self.maxd = maxd
         self.maxp = maxp
         self.minp = minp
         self.iterp = iterp 
@@ -35,6 +37,7 @@ def __init__(self, pp_size, memory_manager:MemoryManager, use_naive_schedule, ma
         # preempt seqs
         self.num_preempt_seqs = 0
         self.log_num_preempt_seqs = 0
+        self.delta_log_num_preempt_seqs = 10
         # num wait tokens
         self.num_wait_tokens = 0
         # abort ids
@@ -102,8 +105,9 @@ def check_preempt(self, num_decode_tokens):
         self.seqs_to_prefill.extendleft(preempt_seqs)
 
         self.num_preempt_seqs += len(preempt_seqs)
-        if self.num_preempt_seqs - self.log_num_preempt_seqs >= 10:
+        if self.num_preempt_seqs - self.log_num_preempt_seqs >= self.delta_log_num_preempt_seqs:
             self.log_num_preempt_seqs = self.num_preempt_seqs
+            self.delta_log_num_preempt_seqs *= 2
             logger.warning(f'#Preempted seqs: {self.num_preempt_seqs}, Try increase --kvthresh or the performance is poor!')
 
     def check_abort_seqs_list(self, seqs:deque, ipc_package:IPCPackage):
@@ -240,6 +244,8 @@ def schedule(self):
             # because we want to solve the situation when #seqs=5 pp_size=4
             decode_token_budget = (
                 num_total_decode_seqs + random.randint(0, self.pp_size-1)) // self.pp_size
+        
+        decode_token_budget = min(self.maxd, decode_token_budget)
 
         self.check_preempt(decode_token_budget)